Итератор объединённого упорядоченного обхода двух источников

37. Объединить два отсортированных итератора

Условие задачи:
Даны два итератора a и b. Каждый из них возвращает элементы в порядке возрастания согласно заданному Comparator.

Необходимо реализовать CollatingIterator<E>, который объединяет оба источника и возвращает все их элементы в общем отсортированном порядке.

Требования:

реализовать методы hasNext() и next();
сохранить все элементы, включая одинаковые;
не считывать исходные итераторы целиком в коллекции;
использовать O(1) дополнительной памяти;
предполагается, что элементы не равны null.

Код:

public class CollatingIterator<E> implements Iterator<E> {

    private final Iterator<E> a;
    private final Iterator<E> b;
    private final Comparator<E> comparator;

    public CollatingIterator(
            Iterator<E> a,
            Iterator<E> b,
            Comparator<E> comparator
    ) {
        this.a = a;
        this.b = b;
        this.comparator = comparator;
    }

    @Override
    public boolean hasNext() {
        return false;
    }

    @Override
    public E next() {
        return null;
    }
}

Пример:

a: 1, 3
b: 2, 3, 5, 6

Результат:
1 → 2 → 3 → 3 → 5 → 6

Спойлеры к решению

Подсказки

💡 Это похоже на этап слияния в Merge Sort.
💡 Достаточно хранить по одному заранее считанному элементу из каждого итератора.
💡 Если элементы есть в обоих источниках, сравни их и верни меньший.
💡 После возврата элемента продвигай только тот итератор, из которого он был взят.
💡 Если один источник закончился, продолжай возвращать элементы второго.

Решение

public class CollatingIterator<E> implements Iterator<E> {

    private final Iterator<E> a;
    private final Iterator<E> b;
    private final Comparator<E> comparator;

    private E nextA;
    private E nextB;

    private boolean hasNextA;
    private boolean hasNextB;

    public CollatingIterator(
            Iterator<E> a,
            Iterator<E> b,
            Comparator<E> comparator
    ) {
        this.a = Objects.requireNonNull(a);
        this.b = Objects.requireNonNull(b);
        this.comparator = Objects.requireNonNull(comparator);

        advanceA();
        advanceB();
    }

    @Override
    public boolean hasNext() {
        return hasNextA || hasNextB;
    }

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

        if (!hasNextA) {
            E result = nextB;
            advanceB();
            return result;
        }

        if (!hasNextB) {
            E result = nextA;
            advanceA();
            return result;
        }

        if (comparator.compare(nextA, nextB) <= 0) {
            E result = nextA;
            advanceA();
            return result;
        }

        E result = nextB;
        advanceB();
        return result;
    }

    private void advanceA() {
        if (a.hasNext()) {
            nextA = a.next();
            hasNextA = true;
        } else {
            nextA = null;
            hasNextA = false;
        }
    }

    private void advanceB() {
        if (b.hasNext()) {
            nextB = b.next();
            hasNextB = true;
        } else {
            nextB = null;
            hasNextB = false;
        }
    }
}

Алгоритм работает как слияние двух отсортированных последовательностей.

В памяти постоянно находятся максимум два элемента:

private E nextA;
private E nextB;

Например, если:

nextA = 1
nextB = 2

то возвращается 1, после чего продвигается только итератор a.

Получаем:

nextA = 3
nextB = 2

Теперь возвращается 2 и продвигается только b.

Если значения равны:

if (comparator.compare(nextA, nextB) <= 0)

сначала возвращается элемент из a, но элемент из b остаётся сохранённым и будет возвращён позже. Поэтому дубликаты не теряются.

Когда один итератор заканчивается:

if (!hasNextA) {
    E result = nextB;
    advanceB();
    return result;
}

элементы продолжают последовательно считываться из второго.

Если закончились оба источника, next() обязан согласно контракту Iterator выбросить:

throw new NoSuchElementException();

Если первый итератор содержит n элементов, а второй — m, полный обход занимает O(n + m) времени.

Каждый отдельный вызов next() выполняется за O(1), а дополнительная память — O(1), поскольку независимо от размера входных последовательностей хранится не более двух предварительно считанных элементов.