Java8 Stream APIの基本(5) - zip の実装 - エンタープライズギークス (Enterprise Geeks)

先日のStream APIの基本(4)では Spliteratorについて解説したが、2つのリストを扱う方法やインデックスアクセスの方法については記述しなかった。
またStream APIの基本(2)では、並列処理では2つの要素から順番に要素を取得するには工夫が必要であると述べた。

一般的に、関数型言語ではzipと呼ばれる方法を用いて、上記の処理を実現する。
今回は Java8 におけるzipの実装例を紹介するが、これはSpliteratorの実装例であり、並列処理でも安全に動作する。 (なお、Java8 が正式リリースされる前にはStreams#zipという同等の機能が存在していたが、正式リリースには見送られたようである。)

`zip`とは

関数型言語のzipは、複数のリストの要素を先頭から順に1つのタプルのリストにまとめる機能である。タプルとは複数の異なる値を持てるデータ型であり、Java には存在しないが Lisp, Haskell, Python, Scala 等でサポートされている。比較的 Java に文法が似ている Scala のコードは次のようになる。

val list1 = List(1,2,3,4)
val list2 = List("A", "B", "C")
list1.zip(list2) // List((1,A), (2,B), (3,C))

上記のコードを実行すると、2つのリストから、タプルのリストが作成される。要素数が異なる場合は短い方に合わせられる。

ペアの作成

一般的に、要素が2つのタプルをペアと呼ぶため、ここではPairというクラスをzipメソッドで扱う型として用意した。 (Pairの要素名を 1,2 としたのはScalaの実装に倣ったためである。)

public class Pair<T, U> {
    final public T _1;
    final public U _2;
    public Pair(T t, U u) {
        _1 = t;
        _2 = u;
    }
    @Override public boolean equals(Object o1){
        if (o1 instanceof Pair) {
            Pair p = (Pair)o1;
            return _1.equals(p._1) && _2.equals(p._2);
        }
        return false;
    }
    @Override public int hashCode() {
        int hash = 1;
        hash = hash * 31 + _1.hashCode();
        hash = hash * 31 + _2.hashCode();
        return hash;
    }
    @Override public String toString(){
        return "(" + _1 + "," + _2 +")";
    }
}

`zip`メソッドの作成

続いてzipメソッドを定義する。ここでは、2つのストリームを1つにまとめるイテレータを用意し、並列ストリームであっても要素を同じ順序で取得できるようにした。前回述べたとおり、サイズがわかるならサイズを指定した方が並列処理の効率が向上する可能性があるため、サイズを引数に取るようにした。ただし冗長である場合は引数を削除し、常にspliteratorUnknownSizeで作成するようにしても良い。

public class StreamUtil {
    
    /** zip サイズ指定 */
    public static <T,U> Stream<Pair<T,U>> zip(Stream<T> s1, Stream<U> s2, int size) {
        PairIterator itr = new PairIterator(s1.iterator(), s2.iterator(), Pair::new);
        int characteristics = Spliterator.IMMUTABLE | Spliterator.NONNULL;
        if (size < 0) {
            return StreamSupport.stream(
                Spliterators.spliteratorUnknownSize(itr, characteristics), false);
        } else {
            return StreamSupport.stream(
                Spliterators.spliterator(itr, size, characteristics), false);
        }
    }
    
    /** zip サイズなし */
    public static <T,U> Stream<Pair<T,U>> zip(Stream<T> s1, Stream<U> s2) {
        return zip(s1, s2, -1);
    }
    
    /** 2つのイテレータから、要素を合成するイテレータ。 */
    private static class PairIterator<T, U, R> implements Iterator<R> {
        private final Iterator<T> i1;
        private final Iterator<U> i2;
        private final BiFunction<T,U,R> mapper;
        public PairIterator(Iterator<T> i1, Iterator<U> i2, BiFunction<T,U,R> mapper) {
            this.i1 = i1;
            this.i2 = i2;
            this.mapper = mapper;
        }
        @Override public boolean hasNext() {
            return i1.hasNext() && i2.hasNext();
        }
        @Override public R next() {
            return mapper.apply(i1.next(), i2.next());
        }
    }
}

実行例

実行例(1) - リストの突合

2つのリストの同じ位置の要素を比較するには以下のようにする。これは2つのリストからどちらの要素も1である最初の要素を取得する例である。

List<Integer> list1 = Arrays.asList(0,0,1,0,1);
List<Integer> list2 = Arrays.asList(1,1,0,0,1);
Optional<Pair<Integer, Integer>> hit = StreamUtil.zip(list1.stream(), list2.stream())
        .filter(p -> p._1.equals(1) && p._1.equals(p._2)).findFirst();
System.out.println(hit.isPresent() ? hit.get() : "none");

実行例(2) - インデックスの使用

ストリームの要素を2つおきに取得したいなど、インデックスを用いる例は以下の通り。
以下は偶数番目の要素のみの合計を求めている。

List<Integer> input = Arrays.asList(1,2,3,4,5,6);
int sum = StreamUtil.zip(input.stream(), Stream.iterate(0, n -> n + 1))
        .parallel() // 並列処理でも問題なく動くか確認する。
        .filter(p -> p._2 % 2 == 1)
        .collect(Collectors.summingInt(p -> p._1));
System.out.println(sum); // 2 + 4 + 6 = 12

0から始まる無限ストリームをzipに適用すると、元のストリームに番号が振られることになる。
その場合、 (1,0), (2,1), (3,2), (4,3), (5,4), (6,5) と、(inputの要素, 番号)　というペアが作成される。
よって、Pair#_2 にfilterを適用することで、従来のインデックスアクセスに相当する処理が記述できる。
これは無限ストリームの有効な使用例の一つでもある。

実行例(3) - 前後の要素の比較

このzipは、1つのリストから前後の要素を比較したい場合にも利用できる。
例として、文字列から同じ文字が連続する回数を求める処理を考えてみよう。 "ABCAABBCCCDEEF" という文字列を対象とした場合、結果は"AA", "BB", "CC", "CC", "EE"の計5回になる（3文字連続の"CCC"は2回とカウント）。

String str = "ABCAABBCCCDEEF";
long count = StreamUtil.zip(str.chars().boxed(), str.chars().skip(1).boxed(), str.length() - 1)
        .peek(p -> System.out.println("before:" + (char)p._1.intValue() 
                          + " after:" + (char)p._2.intValue())) // デバッグ用
        .filter(p -> p._1.equals(p._2)).count();
System.out.println(count); //5

ここでのポイントは、zipに渡す引数だ。
第一引数は文字列そのもののストリームで、第二引数は文字列からskip(1)を実行したストリーム、つまり先頭1文字を削ったストリームである。

"ABCAABBCCCDEEF"と"BCAABBCCCDEEF"をzipすると(A,B), (B,C), (C, A), (A, A) ,,, (E,F) というペアが生成される（peekによるコンソール出力内容を参照）。
これは、文字列の1,2番目、2,3番目、3,4番目の要素のため、ペアの両要素が同じなら同じ文字が連続していることになる。

Scala などでは、このような1つのリストから開始位置をずらしてペアを作る処理にもzipを利用する。従来のループによる書き方では、ループの外に前回の要素を保持する変数などを定義し、ループの初回や最終回に要素の存在判定などを行う必要があった。しかしzipを利用することで、そのような煩雑な処理を記述せずに、同様の処理が実現できる。

まとめ

上記の3つの例を、従来のforループ(特に拡張for文)で記述した場合、一時変数などが増えてコードが煩雑になりがちだが、zipを使って1回の計算に必要な要素を予め用意することで、コードをシンプルに保てる。
また内部的にIteratorを用いることで、並列処理でも安全に2つのストリームを合成できる。

[前多賢太郎]

zipとは