281. Zigzag Iterator

MediumDesignQueueArrayIterator

Problem Description

Given two input vectors v1 and v2, the task is to create an iterator that returns their elements in a zigzag or alternating fashion. In other words, the iterator should return the first element of v1, then the first element of v2, followed by the second element of v1, and so on, alternating between the two vectors.

To achieve this, the iterator should keep track of two pieces of information:

Which vector's turn it is to provide an element.
The next index to be accessed in each vector.

The iterator should also be able to tell whether there are more elements to return, which involves checking both vectors for remaining elements.

Intuition

The solution involves maintaining state for the iterator that includes which vector to access next and the current index within each vector. The core idea is to attempt to retrieve an element from one vector and then switch to the other, continuing this process alternately until all elements from both vectors are exhausted.

In detail, the solution requires:

An initialization step that sets up the structure to hold the vectors and their corresponding indices.
A way to move to the next element (next() function) that adheres to the zigzag order. This includes returning the current element and updating the stored state.
A method to check if there are any elements left (hasNext() function) by iterating over the vectors and ensuring that at least one vector has an element that has not been returned yet.

The next() method retrieves the current element from the active vector based on the current index, then increments the index for that vector and updates which vector is active for the next call, ensuring the alternating pattern.

The hasNext() method checks if the current vector is exhausted by comparing the current index with the length of the vector. If the current vector is exhausted, it moves to the next vector and checks for the same. This process wraps around in a cycle until it returns to the starting vector. If the starting vector is reached and it is also exhausted, the method returns false, indicating that there are no more elements to return. Otherwise, it returns true, indicating that there are further elements to process.

Learn more about Queue patterns.

Solution Approach

To implement the ZigzagIterator, we utilize a few concepts central to the iterator pattern and alternating patterns between two data structures.

Algorithmically, the solution consists of steps performed in the following methods:

__init__: In this constructor method, we initialize critical pieces of information that the iterator needs:
- self.cur: A variable storing which vector is currently active (starting with the first vector, 0).
- self.size: In this particular case, we have two vectors, so the size is static and set to 2. For a more generic solution, it could be the length of an array of vectors.
- self.indexes: An array to track the current index within each vector, initialized to zeros since we start at the beginning of each vector.
- self.vectors: An array that stores the input vectors v1 and v2.
next: When this method is called, we:
- Access the current vector using self.cur to determine whether v1 or v2 should be accessed.
- Retrieve the element at the current index of the active vector.
- Increment the index for the active vector, moving the iterator forward within that vector.
- Update self.cur to point to the next vector, wrapping around using the modulo operator %. This ensures the zigzag pattern by alternating between the vectors.
hasNext: This method is responsible for determining if the iterator has additional elements to process:
- We start by setting a variable start to self.cur to remember the starting point.
- We enter a loop where we check if the current vector (designated by self.cur) has been completely iterated over by comparing the index in self.indexes for the current vector with its length.
- If the current vector is exhausted, we increment self.cur, using % to wrap around, effectively moving to the next vector.
- We keep rotating through the vectors until we return to the starting vector (start), checking if any elements remain in any of them.
- If we return to the starting vector and no elements are left, we return False. Otherwise, at first detection of an available element, we break the loop and return True, confirming that the iterator has more elements.

The data structures used include:

An array (or list in Python) for self.indexes to track the current index within each vector.
An array for self.vectors to hold the input vectors.

A defining pattern of this implementation is the alternating access pattern to vectors, achieved by cycling through using % on self.cur, and a while loop in hasNext to traverse the vectors in round-robin fashion.

The overall time complexity for next is O(1), since it involves only simple operations without any loops. For hasNext, the time complexity in the worst case can be O(m + n), where m and n are the lengths of v1 and v2, respectively, but this only happens when we exhaust all the vectors (when the iterator ends). On average, assuming that the vectors are reasonably balanced in size, the complexity for hasNext is amortized to O(1).

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's illustrate the solution approach using two vectors, v1 and v2, where v1 = [1, 3] and v2 = [2, 4, 5, 6].

Initialization

Create an instance of the ZigzagIterator, initializing the following:
- self.vectors = [[1, 3], [2, 4, 5, 6]]
- self.indexes = [0, 0] (for tracking the current index of each vector)
- self.cur = 0 (starting with the first vector)
- self.size = 2 (we are working with two vectors)

Using `next()`

Call next(), which checks self.cur to determine the current vector. Since self.cur is 0, we look at v1.
The current index for v1 is self.indexes[0], which is 0; hence we return v1[0], which is 1.
Increment self.indexes[0] to 1 and update self.cur to 1 (for the next vector, v2).
Call next() again. This time, self.cur is 1, so we use v2.
We return v2[0], which is 2, and then self.indexes[1] becomes 1. We update self.cur to 0 (modulo 2).
On the next next() call, return v1[1] since self.cur is 0 and self.indexes[0] is 1. The value is 3. Increment self.indexes[0] to 2 and update self.cur to 1.
Continue calling next() to return the next elements from v2: 4, 5, and 6, incrementing self.indexes[1] after each call and toggling self.cur between 1 and 0.

Using `hasNext()`

Initially, self.cur is 0, and self.indexes[0] is 2, which is equal to the length of v1. So we move to v2.
self.indexes[1] is less than the length of v2, meaning there are elements left in v2.
We return True for hasNext().
Once all elements have been processed and self.indexes for both vectors are equal to their respective lengths, hasNext() will return False.

Solution Implementation

1class ZigzagIterator:
2    def __init__(self, v1: List[int], v2: List[int]):
3        # Initialize the current pointer to 0 to start with the first vector
4        self.current = 0
5        # The size variable indicates the number of vectors (always 2 in this case)
6        self.size = 2
7        # Create a list to keep track of the current index in each vector
8        self.indices = [0] * self.size
9        # Store the two input vectors
10        self.vectors = [v1, v2]
11
12    def next(self) -> int:
13        # Identify the current vector and its index
14        vector = self.vectors[self.current]
15        index = self.indices[self.current]
16        # Retrieve the next element from the current vector
17        result = vector[index]
18        # Increment the index for the current vector
19        self.indices[self.current] = index + 1
20        # Move to the next vector for the following call
21        self.current = (self.current + 1) % self.size
22        # Return the retrieved element
23        return result
24
25    def hasNext(self) -> bool:
26        # Store the starting point to avoid infinite loops
27        start = self.current
28        # Check if the current vector is exhausted and move to the next one if necessary
29        while self.indices[self.current] == len(self.vectors[self.current]):
30            # Move to the next vector
31            self.current = (self.current + 1) % self.size
32            # If we loop back to the start, that means all vectors are exhausted
33            if self.current == start:
34                return False
35        # If we haven't returned yet, there's still at least one element left
36        return True
37
38
39# Usage example:
40# i, v = ZigzagIterator(v1, v2), []
41# while i.hasNext(): v.append(i.next())
42

1import java.util.ArrayList;
2import java.util.List;
3
4public class ZigzagIterator {
5    // Current position to determine from which vector to take the next element.
6    private int current;
7  
8    // Total number of vectors
9    private int listCount;
10  
11    // List to keep track of the indexes of current read position for each vector.
12    private List<Integer> indices = new ArrayList<>();
13  
14    // List holding the two input vectors.
15    private List<List<Integer>> vectors = new ArrayList<>();
16
17    // Constructor which initializes the ZigzagIterator with two lists (v1 and v2).
18    public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
19        current = 0;
20        listCount = 2; // Since we always have two lists (v1 and v2) for this iterator.
21        indices.add(0);  // Adding initial indices for both vectors.
22        indices.add(0);
23        vectors.add(v1); // Adding the provided lists to the vectors list.
24        vectors.add(v2);
25    }
26
27    // Returns the next element in the zigzag iteration.
28    public int next() {
29        // Get the current vector and the index for the element to return.
30        List<Integer> vector = vectors.get(current);
31        int index = indices.get(current);
32        int result = vector.get(index); // Get the next element.
33        indices.set(current, index + 1); // Update the index for the current vector.
34        current = (current + 1) % listCount; // Move to the next vector.
35        return result;
36    }
37
38    // Returns true if there is a next element in the iteration, false otherwise.
39    public boolean hasNext() {
40        int start = current; // Remember the starting point.
41        // Cycle through the vectors to find if any vector still has elements left.
42        while (indices.get(current) == vectors.get(current).size()) {
43            current = (current + 1) % listCount; // Move to next vector.
44            if (start == current) {
45                // If we have cycled through all vectors without finding an element,
46                // then there are no elements left, return false.
47                return false;
48            }
49        }
50        // If we have found a vector that still has elements, return true.
51        return true;
52    }
53}
54
55/**
56 * Example usage:
57 * ZigzagIterator iterator = new ZigzagIterator(v1, v2);
58 * List<Integer> result = new ArrayList<>();
59 * while (iterator.hasNext()) {
60 *     result.add(iterator.next());
61 * }
62 */
63

1#include <vector>
2
3class ZigzagIterator {
4private:
5    // Current position to determine from which vector to take the next element
6    int current;
7
8    // Total number of vectors
9    int listCount;
10
11    // Vector to keep track of the indexes of current read position for each vector
12    std::vector<int> indices;
13
14    // Vector holding the two input vectors
15    std::vector<std::vector<int>> vectors;
16
17public:
18    // Constructor that initializes ZigzagIterator with two vectors (v1 and v2)
19    ZigzagIterator(std::vector<int>& v1, std::vector<int>& v2) {
20        current = 0;
21        listCount = 2; // Since we always have two vectors for this iterator
22        indices.push_back(0); // Adding initial indices for both vectors
23        indices.push_back(0);
24        vectors.push_back(v1); // Adding the provided vectors to the vectors list
25        vectors.push_back(v2);
26    }
27
28    // Returns the next element in zigzag iteration
29    int next() {
30        // Get the current vector and index for the element to return
31        std::vector<int>& vector = vectors[current];
32        int index = indices[current];
33        int result = vector[index]; // Get the next element
34        indices[current] = index + 1; // Update the index for the current vector
35        current = (current + 1) % listCount; // Move to the next vector
36        return result;
37    }
38
39    // Returns true if there is a next element in the iteration, false otherwise
40    bool hasNext() {
41        int start = current; // Remember the starting point
42        // Cycle through the vectors to find if any vector still has elements left
43        while (indices[current] == vectors[current].size()) {
44            current = (current + 1) % listCount; // Move to next vector
45            // If we have cycled through all vectors without finding an element,
46            // then there are no elements left, return false
47            if (start == current) {
48                return false;
49            }
50        }
51        // If we found a vector that still has elements, return true
52        return true;
53    }
54};
55
56/**
57 * Example usage:
58 * std::vector<int> v1 = {1,2};
59 * std::vector<int> v2 = {3,4,5,6};
60 * ZigzagIterator iterator = ZigzagIterator(v1, v2);
61 * std::vector<int> result;
62 * while (iterator.hasNext()) {
63 *     result.push_back(iterator.next());
64 * }
65 * // Now 'result' contains the zigzag iteration of v1 and v2
66 */
67

1// Initialize variables to hold state for the iterator
2let current: number = 0; // Holds the position to determine from which array to pick the next element.
3let listCount: number = 2; // In this setup, we always have two arrays defined.
4let indices: number[] = [0, 0]; // Tracks the current read position for each array.
5let vectors: number[][] = []; // Stores the provided arrays for iteration.
6
7// Initializes the iterator with two arrays (v1 and v2).
8function initialize(v1: number[], v2: number[]): void {
9  current = 0;
10  vectors = [v1, v2]; // Stores the two provided arrays (v1 and v2) for zigzag iteration.
11}
12
13// Returns the next element in the zigzag iteration.
14function next(): number {
15  const vector = vectors[current]; // Get the current array.
16  const index = indices[current]; // Get the index for the element to be returned from the current array.
17  const result = vector[index]; // Store the next element.
18  indices[current] = index + 1; // Update the read position for the current array.
19  current = (current + 1) % listCount; // Move to the next array for the following call.
20  return result;
21}
22
23// Checks if there are any more elements in any array.
24function hasNext(): boolean {
25  let start = current; // Save the starting position
26  // Loop through the arrays to determine if any of them still have elements.
27  while (indices[current] === vectors[current].length) {
28    current = (current + 1) % listCount; // Move to the next array.
29    if (start === current) {
30      // If we've checked all arrays and found no remaining elements, return false.
31      return false;
32    }
33  }
34  // Found an array that still has elements, return true.
35  return true;
36}
37
38// Usage
39// initialize(v1, v2);
40// let result: number[] = [];
41// while (hasNext()) {
42//   result.push(next());
43// }
44

Time and Space Complexity

Time Complexity

next() Method: The time complexity for next() is O(1) because we are simply accessing an element from an array, updating an index, and then adjusting the current pointer cur. All these operations take constant time.

hasNext() Method: The time complexity for hasNext() is O(k) in the worst case, where k is the number of vectors (2 in this case). This is because hasNext() might potentially go through both vectors to find out if there's a next element. In the current implementation, k is 2, so we can argue that the method runs in O(1) as well, but if the iterator is extended to support more than two input vectors, the time complexity will depend on the number of vectors, hence O(k).

Overall time complexity per element across the entire iteration process is still O(1). Even though hasNext() may take O(k) time in the worst case, across n elements, it would only be checking each vector once per element, resulting in O(n) for n calls to next(), which averages out to O(1) per call to next().

Space Complexity

The space complexity is O(n1 + n2), where n1 and n2 are the lengths of v1 and v2, respectively. This is due to the storage requirement of the vectors list which contains references to the two input vectors. The additional space used by the iterator object itself (for storing cur and indexes) is O(k) where k is the size of the indexes array or the number of input vectors, which is 2 in the current context.

Since we provided the vectors themselves as input to the constructor, we generally consider them part of the input size rather than attributing them to the auxiliary space complexity of the algorithm. Therefore, if we exclude the space taken by the input vectors, the auxiliary space complexity of the ZigzagIterator implementation is O(1) as we only use a fixed number of additional variables that does not depend on the input size.

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

Which algorithm should you use to find a node that is close to the root of the tree?

DFS

BFS

Binary search

Backtracking

I don’t know