Iterator implementaion for k sorted arrays with no duplicates - interview question

Question

The first part of the question was:
Given k sorted arrays, implement an iterator for iterating over the elements of the arrays in ascending order. For example:

If we have: a1 = {1,3,5}, a2 = {2,4,4,5}, then calling next() method of the iterator implemented for 7 times will return: 1,2,3,4,4,5,5.

I succeed in implementing this part, and wrote code for it below.

The second part was implementing this iterator class when next() method doesn't return duplicates. For the arrays in the example above, if we call the next() method for 5 times we get: 1,2,3,4,5 (if we call it 6 times, we need to get an exception).

I thought it wouldn't be hard - just use HashSet field, add items to this set when entering them to the heap and implement next as a loop that terminates when you get a unique item.
The problem with this approach is that the method hasNext() will not be efficient: you will have to iterate the elements that will be inserted to the heap in future calls for knowing you will actually have unique element in a future call for next().

Do you have idea how to implement this iterator with no duplicates returned, in an efficient way?

import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.PriorityQueue;

public class ComplexIterator implements Iterator<Integer>{

    private class IndexedArrayValue implements Comparable<IndexedArrayValue> {
        int arrayId;
        int index;
        int value;

        public IndexedArrayValue(int arrayId, int index, int value) {
            this.arrayId = arrayId;
            this.index = index;
            this.value = value;
        }

        @Override
        public int compareTo(IndexedArrayValue other) {
            return this.value - other.value;
        }
    }

    private int[][] lists;
    private PriorityQueue<IndexedArrayValue> minHeap;

    public ComplexIterator(int[][] lists) {
        minHeap = new PriorityQueue<IndexedArrayValue>();
        int numOfLists = lists.length;

        this.lists = lists;
        for (int i = 0; i < numOfLists; i++) {
            minHeap.add(new IndexedArrayValue(i, 0, lists[i][0]));
        }
    }

    @Override
    public boolean hasNext() {
        return !this.minHeap.isEmpty();
    }

    @Override
    public Integer next() {
        if (!hasNext())
            throw new NoSuchElementException();

        IndexedArrayValue indArrVal = minHeap.poll();
        int arrayId = indArrVal.arrayId;
        int index = indArrVal.index;
        int value = indArrVal.value;
        int nextIndex = index + 1;

        if (nextIndex < lists[arrayId].length) {
            minHeap.add(new IndexedArrayValue(arrayId, nextIndex, lists[arrayId][nextIndex]));
        }

        return value;
    }

    public static void main (String[] args) {
        int[] arr1 = { 1, 2, 3 };
        int[] arr2 = { 1, 4 };
        int[] arr3 = { 2, 5, 7, 8 };

        int[][] arrs = new int[][] {arr1, arr2, arr3};

        ComplexIterator it = new ComplexIterator(arrs);
        while (it.hasNext()) {
            System.out.print(it.next() + " ");
        }

    }
}

Answer 1

I think a small modification to your original code will eliminate duplicates:

1. When you create the iterator, store the max element of all the arrays (you'll have to check the last element of each of the k arrays to find the maximum).

2. Also store the element returned by the last call to next(). This can be initialized to Integer.MIN_VALUE and modified in each call to next().

3. hasNext() simply checks if the last element returned < max element

4. The new next() calls your original next() repeatedly until it finds an element larger than the previously returned element.

Here's an implementation that modifies your code (it might require some minor modifications to support edge cases such as empty input):

...
private int max; // the maximum element
private int last = Integer.MIN_VALUE; // the last element returned by next()

public ComplexIterator(int[][] lists) {
    minHeap = new PriorityQueue<IndexedArrayValue>();
    int numOfLists = lists.length;

    this.lists = lists;
    max = lists[0][lists[0].length-1];
    for (int i = 0; i < numOfLists; i++) {
        minHeap.add(new IndexedArrayValue(i, 0, lists[i][0]));
        if (lists[i][lists[i].length-1] > max) {
            max = lists[i][lists[i].length-1];
        }
    }
}

@Override
public boolean hasNext() {
    return last < max;
}

@Override
public Integer next() {
    if (!hasNext())
        throw new NoSuchElementException();

    int value;
    do {
        IndexedArrayValue indArrVal = minHeap.poll();
        int arrayId = indArrVal.arrayId;
        int index = indArrVal.index;
        value = indArrVal.value;
        int nextIndex = index + 1;

        if (nextIndex < lists[arrayId].length) {
            minHeap.add(new IndexedArrayValue(arrayId, nextIndex, lists[arrayId][nextIndex]));
        }
    }
    while (value <= last);
    last = value;

    return value;
}