Iterator that yields one array first, then the rest of another

The output is two phases: every element of b in order, then the elements of
a whose value never appears in b, in a's order. The only real work is the
membership test in phase two.

Brute force. Build the whole output list up front: copy b, then for each
x in a scan b to see whether x occurs there. That is O(|a| * |b|) time
and O(|a| + |b|) extra space, and it is not an iterator: the caller pays for
every element even if they stop after the first one.

Optimal. Hash b once into a set, then walk the two arrays lazily with a
cursor:

1. Phase one: return b[j] for j = 0, 1, .... No filtering.
2. Phase two: advance a cursor i over a, skipping every a[i] that is in
   the set. Return the first one that is not.

has_next is the subtle part. In phase two it cannot answer "is i < len(a)?"
because the remaining elements of a may all be filtered out. So the iterator
keeps one invariant:

• After every call, the cursor i points at the next element of a that
  will be returned, or at len(a). A private _advance() restores this by
  skipping members of the set. It runs once in the constructor and after each
  phase-two next.

With that invariant has_next is O(1) and has no side effects, so calling it
twice in a row, or not at all, never changes what next returns. Each element of
a is examined exactly once over the whole iteration, so the total cost is
linear.

Duplicates follow the rule literally. Every copy in b is yielded, and every
copy of a value in a is dropped if the value occurs anywhere in b. (The
"remove one copy per occurrence" variant is a follow-up below.)

from typing import Iterator, List, Sequence, Set


class SubarrayFirstIterator:
    """Yields all of b in order, then the elements of a not present in b."""

    def __init__(self, a: Sequence[int], b: Sequence[int]) -> None:
        self._a = a
        self._b = b
        self._in_b: Set[int] = set(b)
        self._j = 0  # cursor into b (phase one)
        self._i = 0  # cursor into a (phase two)
        self._advance()

    def _advance(self) -> None:
        """Move _i to the next element of a that is not in b (or to len(a))."""
        while self._i < len(self._a) and self._a[self._i] in self._in_b:
            self._i += 1

    def has_next(self) -> bool:
        return self._j < len(self._b) or self._i < len(self._a)

    def next(self) -> int:
        if self._j < len(self._b):
            value = self._b[self._j]
            self._j += 1
            return value
        if self._i < len(self._a):
            value = self._a[self._i]
            self._i += 1
            self._advance()
            return value
        raise StopIteration

    # Python iterator protocol, so it works in for-loops and list().
    def __iter__(self) -> Iterator[int]:
        return self

    def __next__(self) -> int:
        return self.next()


def subarray_first(a: Sequence[int], b: Sequence[int]) -> Iterator[int]:
    """The same order as a generator, for comparison in the interview."""
    in_b = set(b)
    yield from b
    for x in a:
        if x not in in_b:
            yield x


if __name__ == "__main__":
    def drain(it: SubarrayFirstIterator) -> List[int]:
        out: List[int] = []
        while it.has_next():
            out.append(it.next())
        return out

    assert drain(SubarrayFirstIterator([4, 1, 7, 3, 9], [7, 4])) == [7, 4, 1, 3, 9]
    assert drain(SubarrayFirstIterator([5, 2, 5, 8], [5])) == [5, 2, 8]
    assert drain(SubarrayFirstIterator([], [])) == []
    assert drain(SubarrayFirstIterator([1, 2, 3], [])) == [1, 2, 3]
    assert drain(SubarrayFirstIterator([], [6, 6])) == [6, 6]
    assert drain(SubarrayFirstIterator([1, 1, 1], [1])) == [1]
    assert drain(SubarrayFirstIterator([2, 3], [9, 2, 9])) == [9, 2, 9, 3]

    it = SubarrayFirstIterator([3, 3], [3])
    assert it.has_next() and it.has_next()  # idempotent
    assert it.next() == 3
    assert not it.has_next()                # trailing a is all filtered
    try:
        it.next()
        raise AssertionError("expected StopIteration")
    except StopIteration:
        pass

    assert list(SubarrayFirstIterator([4, 1, 7], [7])) == [7, 4, 1]
    assert list(subarray_first([4, 1, 7, 3, 9], [7, 4])) == [7, 4, 1, 3, 9]
    print("all tests passed")

Let n = len(a) and m = len(b).

• Constructor: O(m) to build the set, plus the first _advance, which is
  O(n) in the worst case but is paid for by the amortized bound below.
• next: O(1) amortized. Each element of a is skipped at most once
  across the whole iteration, so a full pass costs O(n + m) in total.
• has_next: O(1) worst case, no side effects.
• Space: O(u) for the set, where u is the number of distinct values in
  b, plus two cursors. The inputs are referenced, not copied.

• a = [4, 1, 7, 3, 9], b = [7, 4] → 7, 4, 1, 3, 9
• a = [5, 2, 5, 8], b = [5] → 5, 2, 8 (every copy of 5 in a is dropped)
• a = [1, 1, 1], b = [1] → 1
• a = [2, 3], b = [9, 2, 9] → 9, 2, 9, 3 (duplicates in b all yielded;
  b may hold values not in a)
• a = [1, 2, 3], b = [] → 1, 2, 3
• a = [], b = [6, 6] → 6, 6
• a = [], b = [] → has_next() is false at once; next() raises
• a = [3, 3], b = [3] → after the first next() returns 3, has_next()
  must be false even though a still has unread elements
• Calling has_next() several times in a row returns the same answer and does
  not skip anything

1. Multiset semantics ("remove one copy per occurrence"). If b is a
   sub-multiset of a and the rest of a means a minus b with
   multiplicity, keep a Counter of b instead of a set. In phase two, when
   a[i] has a positive count, decrement it and skip; otherwise yield.
   a = [5, 2, 5, 8], b = [5] then gives 5, 2, 5, 8. Same complexity.
2. Inputs are iterators, not arrays. b can still be streamed in phase one,
   but the set must be complete before phase two starts. Build it while
   yielding b (add each value as you return it). a needs only forward
   access, so it can be any iterator; has_next then needs a one-element
   lookahead buffer (a "peeking" wrapper) instead of an index.
3. b is too large for memory. Trade space for time. If both arrays can
   be sorted copies, a two-pointer merge answers membership; to keep a's
   order, sort (value, index) pairs of a, mark the excluded indices in a
   bitset of size n, then stream a. For approximate cases, a Bloom filter
   gives no false negatives, so it can only over-filter, which must be
   discussed.
4. More than two arrays. Given b1, b2, ..., bk and then a, yield each
   bi in turn, then a minus the union. Generalize phase one to a list of
   sources and a source index; the set is the union. A variant asks to dedupe
   across all sources, so each value appears once: keep a seen set and filter
   every phase, not only the last.
5. Add remove() or peek(). peek is free given the invariant: return
   b[j] or a[i] without moving. remove on the underlying collection
   needs care about which cursor last moved, so store a "last returned from"
   tag.

Mid-level (L4)

• States the order correctly and uses a hash set rather than a nested scan.
• Delivers a working next / has_next (or a generator) that handles empty
  arrays.
• Catches the trailing-filtered-elements case once prompted by a test.

Senior (L5)

• Spots unprompted that has_next cannot be i < len(a) and keeps an explicit
  invariant (pre-advance or lookahead) so has_next is idempotent and O(1).
• Asks how duplicates should behave and can switch to Counter semantics.
• Gives amortized complexity and explains why each element is examined once.

Staff (L6+)

• Separates the concerns: a generic filtered iterator plus a concatenation of
  iterators, so k sources, streams, and different filters compose.
• Discusses streaming inputs, memory limits on b, and what a lossy filter
  would do to correctness.
• Covers API behavior: what next does when exhausted, whether mutating the
  inputs during iteration is allowed, and thread safety if it matters.