Iterator over multiple SortedSet objects

Question

In Java, I have several SortedSet instances. I would like to iterate over the elements from all these sets. One simple option is to create a new SortedSet, such as TreeSet x, deep-copy the contents of all the individual sets y_1, ..., y_n into it using x.addAll(y_i), and then iterate over x.

But is there a way to avoid deep copy? Couldn't I just create a view of type SortedSet which would somehow encapsulate the iterators of all the inner sets, but behave as a single set?

Answer 1

I'd prefer an existing, tested solution, rather than writing my own.

I'm not aware of any existing solution to accomplish this task, so I took the time to write one for you. I'm sure there's room for improvement on it, so take it as a guideline and nothing else.

As Sandor points out in his answer, there are some limitations that must be imposed or assumed. One such limitation is that every SortedSet must be sorted relative to the same order, otherwise there's no point in comparing their elements without creating a new set (representing the union of every individual set).

Here follows my code example which, as you'll notice, is relatively more complex than just creating a new set and adding all elements to it.

import java.util.*;

final class MultiSortedSetView<E> implements Iterable<E> {

    private final List<SortedSet<E>> sets = new ArrayList<>();
    private final Comparator<? super E> comparator;

    MultiSortedSetView() {
        comparator = null;
    }

    MultiSortedSetView(final Comparator<? super E> comp) {
        comparator = comp;
    }


    @Override
    public Iterator<E> iterator() {
        return new MultiSortedSetIterator<E>(sets, comparator);
    }


    MultiSortedSetView<E> add(final SortedSet<E> set) {
        // You may remove this `if` if you already know
        // every set uses the same comparator.
        if (comparator != set.comparator()) {
            throw new IllegalArgumentException("Different Comparator!");
        }
        sets.add(set);
        return this;
    }


    @Override
    public boolean equals(final Object o) {
        if (this == o) { return true; }
        if (!(o instanceof MultiSortedSetView)) { return false; }
        final MultiSortedSetView<?> n = (MultiSortedSetView<?>) o;
        return sets.equals(n.sets) &&
                (comparator == n.comparator ||
                (comparator != null ? comparator.equals(n.comparator) :
                    n.comparator.equals(comparator)));
    }

    @Override
    public int hashCode() {
        int hash = comparator == null ? 0 : comparator.hashCode();
        return 37 * hash + sets.hashCode();
    }

    @Override
    public String toString() {
        return sets.toString();
    }



    private final static class MultiSortedSetIterator<E>
            implements Iterator<E> {

        private final List<Iterator<E>> iterators;
        private final PriorityQueue<Element<E>> queue;

        private MultiSortedSetIterator(final List<SortedSet<E>> sets,
                final Comparator<? super E> comparator) {
            final int n = sets.size();
            queue = new PriorityQueue<Element<E>>(n,
                    new ElementComparator<E>(comparator));
            iterators = new ArrayList<Iterator<E>>(n);
            for (final SortedSet<E> s: sets) {
                iterators.add(s.iterator());
            }
            prepareQueue();
        }


        @Override
        public E next() {
            final Element<E> e = queue.poll();
            if (e == null) {
                throw new NoSuchElementException();
            }
            if (!insertFromIterator(e.iterator)) {
                iterators.remove(e.iterator);
            }
            return e.element;
        }

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


        private void prepareQueue() {
            final Iterator<Iterator<E>> iterator = iterators.iterator();
            while (iterator.hasNext()) {
                if (!insertFromIterator(iterator.next())) {
                    iterator.remove();
                }
            }
        }

        private boolean insertFromIterator(final Iterator<E> i) {
            while (i.hasNext()) {
                final Element<E> e = new Element<>(i.next(), i);
                if (!queue.contains(e)) {
                    queue.add(e);
                    return true;
                }
            }
            return false;
        }



        private static final class Element<E> {
            final E element;
            final Iterator<E> iterator;

            Element(final E e, final Iterator<E> i) {
                element = e;
                iterator = i;
            }

            @Override
            public boolean equals(final Object o) {
                if (o == this) { return true; }
                if (!(o instanceof Element)) { return false; }
                final Element<?> e = (Element<?>) o;
                return element.equals(e.element);
            }
        }


        private static final class ElementComparator<E>
                implements Comparator<Element<E>> {
            final Comparator<? super E> comparator;

            ElementComparator(final Comparator<? super E> comp) {
                comparator = comp;
            }

            @Override
            @SuppressWarnings("unchecked")
            public int compare(final Element<E> e1, final Element<E> e2) {
                if (comparator != null) {
                    return comparator.compare(e1.element, e2.element);
                }
                return ((Comparable<? super E>) e1.element)
                        .compareTo(e2.element);
            }
        }
    }
}

The inner workings of this class are simple to grasp. The view keeps a list of sorted sets, the ones you want to iterate over. It also needs the comparator that will be used to compare elements (null to use their natural ordering). You can only add (distinct) sets to the view.

The rest of the magic happens in the Iterator of this view. This iterator keeps a PriorityQueue of the elements that will be returned from next() and a list of iterators from the individual sets.
This queue will have, at all times, at most one element per set, and it discards repeating elements. The iterator also discards empty and used up iterators. In short, it guarantees that you will traverse every element exactly once (as in a set).

Here's an example on how to use this class.

SortedSet<Integer> s1 = new TreeSet<>();
SortedSet<Integer> s2 = new TreeSet<>();
SortedSet<Integer> s3 = new TreeSet<>();
SortedSet<Integer> s4 = new TreeSet<>();

// ...

MultiSortedSetView<Integer> v =
        new MultiSortedSetView<Integer>()
            .add(s1)
            .add(s2)
            .add(s3)
            .add(s4);

for (final Integer i: v) {
    System.out.println(i);
}