Problem

Implement a 2D iterator class. It will be initialized with an array of arrays, and should implement the following methods:

  • next(): returns the next element in the array of arrays. If there are no more elements, raise an exception.
  • has_next(): returns whether or not the iterator still has elements left.

Examples

Example 1

1
2
3
4

Input: [[1, 2], [3], [], [4, 5, 6]]
Output: 1, 2, 3, 4, 5, 6
Explanation: 
Calling `next()` repeatedly on the 2D iterator initialized with [[1, 2], [3], [], [4, 5, 6]] should output the sequence 1, 2, 3, 4, 5, 6.

Similar Problem

Flatten Nested List Iterator

Solution

Method 1 - Check the size of the array

To implement a 2D iterator:

  1. Initialization: Begin by pointing to the first element of the first sub-array. Skip any empty sub-arrays.
  2. next() method: Return the current element and move the pointer to the next element. If the end of a sub-array is reached, move to the next non-empty sub-array.
  3. has_next() method: Check if there are any more elements left in the 2D array by ensuring the pointers are within valid bounds and not pointing to an empty sub-array.

Approach

  1. Initialization: Use indices to track the current position within the 2D array.
  2. next() method: Move the indices forward after returning the current element.
  3. has_next() method: Check the current indices and skip any empty sub-arrays.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

public class Solution {
    public class TwoDIterator {
        private List<List<Integer>> array;
        private int row;
        private int col;
        
        public TwoDIterator(List<List<Integer>> array) {
            this.array = array;
            this.row = 0;
            this.col = 0;
            advanceToNext();
        }
        
        public boolean has_next() {
            return row < array.size();
        }
        
        public int next() {
            if (!has_next()) {
                throw new NoSuchElementException();
            }
            int result = array.get(row).get(col);
            col++;
            advanceToNext();
            return result;
        }
        
        private void advanceToNext() {
            while (row < array.size() && col == array.get(row).size()) {
                row++;
                col = 0;
            }
        }
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

class Solution:
    class TwoDIterator:
        def __init__(self, array: List[List[int]]):
            self.array = array
            self.row = 0
            self.col = 0
            self._advance_to_next()
        
        def has_next(self) -> bool:
            return self.row < len(self.array)
        
        def next(self) -> int:
            if not self.has_next():
                raise StopIteration("No more elements")
            result = self.array[self.row][self.col]
            self.col += 1
            self._advance_to_next()
            return result
        
        def _advance_to_next(self):
            while self.row < len(self.array) and self.col == len(self.array[self.row]):
                self.row += 1
                self.col = 0

Complexity

  • ⏰ Time complexity:
    • O(1) for next() since we are only moving indices forward after each call.
    • O(1) for has_next() since the check is done in constant time.
  • 🧺 Space complexity: O(1) since we are only using a fixed amount of extra space to keep track of indices.

Method 2 -

Code

Complexity

  • ⏰ Time complexity: O(NNNXXXNNN)
  • 🧺 Space complexity: O(NNNXXX)