思いつくままに

Java 8のStream API。Zip Withが標準ライブラリから落ちてしまいました。
使用頻度は高いので、用意してみました。

public class ZipWithHelper {
  private static class ZipWithIterator<A, B, C> implements Iterator<C> {
    private final Iterator<? extends A> aIterator;
    private final Iterator<? extends B> bIterator;
    private final BiFunction<? super A, ? super B, C> mapper;

    public ZipWithIterator(Iterator<? extends A> aIterator, Iterator<? extends B> bIterator,
        BiFunction<? super A, ? super B, C> mapper) {
      this.aIterator = aIterator;
      this.bIterator = bIterator;
      this.mapper = mapper;
    }

    @Override
    public boolean hasNext() {
      return aIterator.hasNext() && bIterator.hasNext();
    }

    @Override
    public C next() {
      return mapper.apply(aIterator.next(), bIterator.next());
    }
  }

  public static <A, B, C> Stream<C> zipWith(Stream<? extends A> as, Stream<? extends B> bs,
      BiFunction<? extends A, ? extends B, C> mapper) {
    ZipWithIterator iterator = new ZipWithIterator(as.iterator(), bs.iterator(), mapper);
    return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false);
  }
}

Zip Withもそうなんですが、もともとは、Stream＜A＞、Stream＜B＞の結合処理をしたいよね、という話で始まっています。

どちらか片方Stream(Aとします)のキーでMapを作って、そのMapに対してもう片方のStream(B)でループを回せばいいじゃないかという話なのですが、
Stream A、Bのサイズが共に大きくて、Mapをメモリ上に確保する余裕がない場合に、せっかくStream APIがあるんだから、Inner Join的な結合(歯抜けになっている可能性のあるZip With)をする処理を書けばいいじゃないかという話になったのです。
つまり、メモリ消費量は定数、計算量はO(N+M)ということですね。

その時は、「驚くべきコードを見つけたがそれを書くには余白が云々」的なノリでごまかしたのですが、えー書けないの？と思われるのも癪なので書いてみました。

public class MergeJoinHelper {
  private static class MergeJoinIterator<A, B, C> implements Iterator<C> {
    private final Iterator<? extends A> aIterator;
    private final Iterator<? extends B> bIterator;
    private final MergeJoinComparator<? super A, ? super B> comparator;
    private final BiFunction<? super A, ? super B, C> mapper;

    private A currentA = null;
    private B currentB = null;
    private A nextA = null;
    private B nextB = null;
    private boolean searched = false;

    public MergeJoinIterator(Iterator<? extends A> aIterator, Iterator<? extends B> bIterator,
        MergeJoinComparator<? super A, ? super B> comparator, BiFunction<? super A, ? super B, C> mapper) {
      this.aIterator = aIterator;
      this.bIterator = bIterator;
      this.comparator = comparator;
      this.mapper = mapper;
      this.nextA = aIterator.hasNext() ? aIterator.next() : null;
      this.nextB = bIterator.hasNext() ? bIterator.next() : null;
    }

    public boolean hasNext() {
      return searchForNextCandidate();
    }

    public C next() {
      if (!searchForNextCandidate()) {
        throw new NoSuchElementException();
      }
      searched = false;
      return mapper.apply(currentA, currentB);
    }

    private boolean searchForNextCandidate() {
      if (searched) {
        return comparator.compare(currentA, currentB) == 0;
      }
      searched = true;
      while (nextA != null && nextB != null) {
        while (comparator.compare(nextA, nextB) > 0) {
          if (!bIterator.hasNext()) {
            nextB = null;
            return false;
          }
          nextB = bIterator.next();
        }
        while (comparator.compare(nextA, nextB) < 0) {
          if (!aIterator.hasNext()) {
            nextA = null;
            return false;
          }
        }
        if (comparator.compare(nextA, nextB) == 0) {
          currentA = nextA;
          currentB = nextB;
          nextA = aIterator.hasNext() ? aIterator.next() : null;
          nextB = bIterator.hasNext() ? bIterator.next() : null;
          return true;
        }
      }
      return false;
    }
  }

  public static <A, B, C> Stream<C> mergeJoin(Stream<? extends A> as, Stream<? extends B> bs,
      MergeJoinComparator<? super A, ? super B> comparator, BiFunction<? super A, ? super B, C> mapper) {
    MergeJoinIterator<? extends A, ? extends B, C> iterator = new MergeJoinIterator<>(as.iterator(), bs.iterator(), comparator, mapper);
    return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator,
      Spliterator.ORDERED), false);
  }
}

public interface MergeJoinComparator<A, B> {
  public int compare(A a, B b);
}

［A_10, A_20, A_30］
［B_0, B_10, B_30］

から、

［C_10(A_10, B_10), C_30(A_30, B_30)］

を生成してくれます。

一応制約があって、Stream＜A＞、Stream＜B＞がMergeJoinComparator的に厳密な単調増加になっている必要があります。
要するに、

［A_10, A_20, A_20］

みたいにキーが重複する場合は動作が保証されない点です。
前提条件として、キーは歯抜けがありうるけど（該当するレコードが存在しない可能性があるけど）、キーはキーとして一意になっているということです。
（そうでないと、計算量もメモリ消費量もどのみちO(NM)になって、「メモリ消費量は定数、計算量はO(N+M)」が実現できません）

あとは、対応するキーが存在しない（相方がnull）の場合に、outer join的な結合をするフラグ的なものは欲しいところです。

思いつくままに

Java 8でStreamのMerge処理 (続)

Java 8でStreamのMerge処理

Delphiで連想配列

Lambdajって面白そうなライブラリを見つけた

Java/CassandraのHello World!

リスト系コンポーネントにオブジェクトを関連づけて利用する

リストに大量のアイテムを登録する