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
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
Solution
Method 1 - Check the size of the array
To implement a 2D iterator:
- Initialization: Begin by pointing to the first element of the first sub-array. Skip any empty sub-arrays.
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.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
- Initialization: Use indices to track the current position within the 2D array.
next()
method: Move the indices forward after returning the current element.has_next()
method: Check the current indices and skip any empty sub-arrays.
Code
Java
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;
}
}
}
}
Python
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)
fornext()
since we are only moving indices forward after each call.O(1)
forhas_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
Java
Python
Complexity
- ⏰ Time complexity:
O(NNNXXXNNN)
- 🧺 Space complexity:
O(NNNXXX)