Composite Iterator. Java

March 8, 2016 Collections Iterable Iterator Java

Imagine you have two collections with different types of data. Both of them are sorted by some rule. You need to iterate over both collections, preserving their sort order. Like this:

We want to share a simple but effective technique using the Iterator and Iterable interfaces.

Let’s suppose that we have the following class hierarchy:

class Foo implements Comparable<Foo> {
    final int id;

    Foo(int id) {
        this.id = id;
    }

    @Override
    public int compareTo(Foo other) {
        return id - other.id;
    }

    @Override
    public String toString() {
        return getClass().getName() + "-" + id;
    }
}

class Bar extends Foo {
    Bar(int id) {
        super(id);
    }
}

class Qux extends Foo {
    Qux(int id) {
        super(id);
    }
}

Method compareTo() compares objects by the id field in natural order. Method toString() returns the exact name of the object’s class and its id.

Once upon a time we faced a similar hierarchy in production code. In addition, there were two types of lists: List<Bar> and List<Qux>. Let us remind you that both lists are sorted!

List<Bar> list1 = Arrays.asList(new Bar(10), new Bar(20), new Bar(50), new Bar(80), new Bar(90));
List<Qux> list2 = Arrays.asList(new Qux(30), new Qux(40), new Qux(60), new Qux(70), new Qux(100));

So, the task is to iterate over these two lists efficiently: linear time complexity O(n), no copies of collections, no additional sorting, no unnecessary loops.

“The simplest” solution is to write two loops, then get two elements (one from each list), compare them, decide which one to use, and iterate again. If one list becomes empty too early, we need to handle this situation in a separate loop. This code will be too complex and not reusable. (We don’t want to provide an example here; you may write it yourself as an exercise.)

On the other hand, we can avoid a lot of complications by using the Iterator interface. Let’s declare the CompositeIterator class this way:

public class CompositeIterator<T> implements Iterator<T> {
    private final Iterator<? extends T> it1;
    private final Iterator<? extends T> it2;
    private final Comparator<T> comparator;

    private Iterator<? extends T> lastUsed;
    private T obj1;
    private T obj2;

Fields it1 and it2 refer to the iterators of both collections. Comparator is used for ordering. Field lastUsed stores the iterator which was last used to return a value. Fields obj1 and obj2 refer to the objects retrieved from it1 and it2 respectively. Well, how will it work?

public boolean hasNext() {
    return obj1 != null || obj2 != null || it1.hasNext() || it2.hasNext();
}

public T next() {
    if (obj1 == null) {
        obj1 = (it1.hasNext()) ? it1.next() : null;
    }

    if (obj2 == null) {
        obj2 = (it2.hasNext()) ? it2.next() : null;
    }

    if (obj1 != null && obj2 != null) {
        if (comparator.compare(obj1, obj2) < 0) {
            return obj1();
        } else {
            return obj2();
        }
    } else if (obj1 != null) {
        return obj1();
    } else if (obj2 != null) {
        return obj2();
    } else {
        throw new NoSuchElementException();
    }
}

private T obj1() {
    lastUsed = it1;
    T result = obj1;
    obj1 = null;
    return result;
}

private T obj2() {
    lastUsed = it2;
    T result = obj2;
    obj2 = null;
    return result;
}

public void remove() {
    if (lastUsed == null) {
        throw new IllegalStateException();
    }

    lastUsed.remove();
    lastUsed = null;
}

Method hasNext() checks whether at least one object (obj1 or obj2) has already been retrieved. If both of them are null, it checks whether at least one of the iterators has more elements by calling hasNext().

If one of these four conditions is true, we can proceed to the next() method. Here we need to make sure that we use both iterators, so if obj1 or obj2 is null, we need to get it from its iterator. If both of them are not null, we need to compare them and choose which one to return. Otherwise we just return the one which is not null. If both of them are still null at this point, it means that next() was called incorrectly, so we need to throw NoSuchElementException.

Pay attention to the two methods obj1() and obj2(). They are used to save a reference to the last used iterator and to clear the reference to the object which was returned.

Method remove() uses the lastUsed field to remove the last returned element from the proper iterator.

To make this iterator easy to use, we will create a simple implementation of the Iterable interface, which creates a new CompositeIterator every time the iterator() method is called:

public class CompositeIterable<T> implements Iterable<T> {
    private final Iterable<? extends T> iterable1;
    private final Iterable<? extends T> iterable2;

    public CompositeIterable(Iterable<? extends T> iterable1, Iterable<? extends T> iterable2) {
        this.iterable1 = iterable1;
        this.iterable2 = iterable2;
    }

    @Override
    public Iterator<T> iterator() {
        return new CompositeIterator<T>(iterable1.iterator(), iterable2.iterator());
    }
}

That’s all! Don’t believe it? Try it!

public static void main(String[] args) {
    List<Bar> list1 = Arrays.asList(new Bar(10), new Bar(20), new Bar(50), new Bar(80), new Bar(90));
    List<Qux> list2 = Arrays.asList(new Qux(30), new Qux(40), new Qux(60), new Qux(70), new Qux(100));

    Iterable<Foo> composite = new CompositeIterable<>(list1, list2);

    for (Foo foo : composite) {
        System.out.println(foo);
    }
}

Wow! We didn’t even expect such pretty code. Here is the output:

Bar-10
Bar-20
Qux-30
Qux-40
Bar-50
Qux-60
Qux-70
Bar-80
Bar-90
Qux-100

In addition, we want to mention that you can use this approach for any number of wrapped collections. Two, three, four — it doesn’t matter! You can wrap them in pairs many times.

We hope this solution will help you process collections better. You may find all files from this example at the end of this post. Thanks for reading!

Source code

CompositeIterator.java

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;

public class CompositeIterator<T> implements Iterator<T> {
    private final Iterator<? extends T> it1;
    private final Iterator<? extends T> it2;
    private final Comparator<T> comparator;

    private Iterator<? extends T> lastUsed;
    private T obj1;
    private T obj2;

    public CompositeIterator(Iterator<? extends T> it1, Iterator<? extends T> it2) {
        this(it1, it2, null);
    }

    public CompositeIterator(Iterator<? extends T> it1, Iterator<? extends T> it2, Comparator<T> comparator) {
        if (it1 == null) {
            throw new NullPointerException();
        }

        if (it2 == null) {
            throw new NullPointerException();
        }

        this.it1 = it1;
        this.it2 = it2;
        this.comparator = comparator != null ? comparator : defaultComparator();
    }

    public boolean hasNext() {
        return obj1 != null || obj2 != null || it1.hasNext() || it2.hasNext();
    }

    public T next() {
        if (obj1 == null) {
            obj1 = (it1.hasNext()) ? it1.next() : null;
        }

        if (obj2 == null) {
            obj2 = (it2.hasNext()) ? it2.next() : null;
        }

        if (obj1 != null && obj2 != null) {
            if (comparator.compare(obj1, obj2) < 0) {
                return obj1();
            } else {
                return obj2();
            }
        } else if (obj1 != null) {
            return obj1();
        } else if (obj2 != null) {
            return obj2();
        } else {
            throw new NoSuchElementException();
        }
    }

    private T obj1() {
        lastUsed = it1;
        T result = obj1;
        obj1 = null;
        return result;
    }

    private T obj2() {
        lastUsed = it2;
        T result = obj2;
        obj2 = null;
        return result;
    }

    public void remove() {
        if (lastUsed == null) {
            throw new IllegalStateException();
        }

        lastUsed.remove();
        lastUsed = null;
    }

    @SuppressWarnings("unchecked")
    private Comparator<T> defaultComparator() {
        return (o1, o2) -> ((Comparable<T>) o1).compareTo(o2);
    }
}

CompositeIterable.java

import java.util.Iterator;

public class CompositeIterable<T> implements Iterable<T> {
    private final Iterable<? extends T> iterable1;
    private final Iterable<? extends T> iterable2;

    public CompositeIterable(Iterable<? extends T> iterable1, Iterable<? extends T> iterable2) {
        this.iterable1 = iterable1;
        this.iterable2 = iterable2;
    }

    @Override
    public Iterator<T> iterator() {
        return new CompositeIterator<>(iterable1.iterator(), iterable2.iterator());
    }
}

TestCompositeIterator.java

import java.util.Arrays;
import java.util.List;

public class TestCompositeComparator {

    public static void main(String[] args) {
        List<Bar> list1 = Arrays.asList(new Bar(10), new Bar(20), new Bar(50), new Bar(80), new Bar(90));
        List<Qux> list2 = Arrays.asList(new Qux(30), new Qux(40), new Qux(60), new Qux(70), new Qux(100));

        Iterable<Foo> composite = new CompositeIterable<>(list1, list2);

        for (Foo foo : composite) {
            System.out.println(foo);
        }
    }

}

class Foo implements Comparable<Foo> {
    final int id;

    Foo(int id) {
        this.id = id;
    }

    @Override
    public int compareTo(Foo other) {
        return id - other.id;
    }

    @Override
    public String toString() {
        return getClass().getName() + "-" + id;
    }
}

class Bar extends Foo {
    Bar(int id) {
        super(id);
    }
}

class Qux extends Foo {
    Qux(int id) {
        super(id);
    }
}