Combine two sorted iterators in O(1) time

Question

I was asked the following question in an interview:

Combine two iterators over their sorted contents such that the resulting iterator should iterate over the combination of these 2 iterators in sorted order in O(1) time (these iterators iterate over a String).

I wrote the below code but I'm sure it doesn't perform in O(1) time. What advice do you have for matching the constraints set by the interview question?

import java.util.Iterator;
import java.util.Set;
import java.util.TreeSet;

public class iteratorCombine {

// assumption1: elements are hardcoded
// assumption2: both iterators have equal number of elements
public static void main(String[] args) {

    iteratorCombine testObj = new iteratorCombine();
    Set<String> firstSet = new TreeSet<String>();
    Set<String> secondSet = new TreeSet<String>();
    Set<String> combinedSet;
    firstSet = testObj.storeElements1(firstSet);
    secondSet = testObj.storeElements2(secondSet);

    Iterator<String> it1 = firstSet.iterator();
    Iterator<String> it2 = secondSet.iterator();

    combinedSet = testObj.combine(it1, it2);

    // output
    Iterator<String> itComb = combinedSet.iterator();
    while(itComb.hasNext()){
        System.out.println(itComb.next());
    }

}

public Set<String> storeElements1(Set<String> firstSet){
    firstSet.add("first3");
    firstSet.add("first1");
    firstSet.add("first2");
    return firstSet;
}

public Set<String> storeElements2(Set<String> secondSet){
    secondSet.add("second3");
    secondSet.add("second1");
    secondSet.add("second2");
    return secondSet;
}

public Set<String> combine(Iterator<String> it1, Iterator<String>it2){
    String firstEle, secondEle;
    Set<String> combinedSet = new TreeSet<String>();
    while (it1.hasNext() && it2.hasNext()) {
        firstEle = it1.next();
        secondEle = it2.next();
        combinedSet.add(firstEle+secondEle);
    }
    return combinedSet;
  }
}

Answer 1

I believe that you can't do it if you don't extend iterator and support a peek function. Such an iterator is not that hard. Here is a way for doing it.

static class PeekingIterator<T> implements Iterator<T> {
    private final Iterator<T> iterator;
    private T temp;

    public PeekingIterator(Iterator<T> iterator) {
        this.iterator = iterator;
    }

    public T peek() {
        //if there is no peek, advance the iterator and store its value, return the peek otherwise
        if(temp==null){ 
            temp = this.iterator.next();
        }
        return temp;
    }

    @Override
    public T next() {
       //if we already have a peek,return it and nullify it, otherwise do normal next()
        if(temp!=null){
            T t = temp;
            temp = null;
            return t;
        }else{
            return this.iterator.next();
        }
    }

    @Override
    public boolean hasNext() {
        return this.iterator.hasNext() || temp!=null;
    }
}

Once you can peek, the rest is easy, you can build SortedIterator using two peeking iterators, peek both iterators and advance the iterator that has the smaller element.

static class SortedIterator<T extends Comparable<T>> implements Iterator<T>{
    private final PeekingIterator<T> peekingIterator1;
    private final PeekingIterator<T> peekingIterator2;

    SortedIterator(Iterator<T> source1, Iterator<T> source2){
        peekingIterator1 = new PeekingIterator<>(source1);
        peekingIterator2 = new PeekingIterator<>(source2);
    }

    @Override
    public boolean hasNext() {
        return peekingIterator1.hasNext() || peekingIterator2.hasNext();
    }

    @Override
    public T next() {
        if(!peekingIterator1.hasNext()){
            return peekingIterator2.next();
        }
        if(!peekingIterator2.hasNext()){
            return peekingIterator1.next();
        }

        T peek1 = peekingIterator1.peek();
        T peek2 = peekingIterator2.peek();
        if(peek1.compareTo(peek2)<0){
            return peekingIterator1.next();
        }
        return peekingIterator2.next();
    }
}

The analysis are obvious here, SortedIterator.next and SortedIterator.hasNext run in constant time.