Java8 Stream APIを使ったピタゴラス数生成処理

先日の記事では、Stream APIの基本的な仕組みを書いた。今回はサンプルコードとして、ピタゴラス数を生成する処理をStreamで書いてみる。

ピタゴラス数とは、三平方の定理でおなじみの
a * a + b * b = c * c　
を満たす自然数の組のことだ。
最初のピタゴラス数は(3,4,5)で、(6,8,10)のような倍数となる組み合わせは除外する。

このロジックは、 a < b < c を前提として、
1. 最初に自然数cを生成し、
2. 各cの値について、c未満の自然数bを生成して、
3. 同様にb未満の数aを生成する。

候補となる3つの数字の組み合わせを生成した後で、
4. 数a,b,cについて条件を満たす組み合わせのみを抽出する
という流れで実現できる。

コードは以下のようになる。

import java.util.stream.IntStream;
import java.util.stream.Stream;

public class PitagorasNumberGenerator {
    public static void main(String[] args) {
       // ピタゴラス数を3つ表示する。
       generate().limit(3).forEach(System.out::println);
       // 100番目のピタゴラス数を表示する。
       generate().skip(99).limit(1).forEach(System.out::println);
        
    }
    /** ピタゴラス数生成 */
    public static Stream<String> generate() {
        return num(1) // cを無限数として生成。
                .flatMap(c -> num(1).limit(c) // 各数ごとに、1からc-1までの数を生成してbとする。
                    .flatMap(b -> num(1).limit(b) // 同様にb未満の数を生成してaとする。
                        .filter(a -> isCoprime(a, b)) // a,bだけて互いに素な数だけ抽出。
                        .filter(a -> isPitagorasNumber(a, b, c)) // ピタゴラス数の条件を満たすものを抽出
                        .map(a -> String.format("(%d,%d,%d)", a,b,c)))); // 最後はStringに変換。
    }
    
    /** 無限数生成 */
    public static Stream<Integer> num(Integer init) {
        return  IntStream.iterate(init, x -> x + 1).boxed();
    }
    /** 互いに素であるかの判定 */
    private static boolean isCoprime(int a, int b) {
        if (b == 0) return false;
        if (b == 1) return true;
        if (b > a) return isCoprime(b, a);
        return isCoprime(b, a % b);
    }
    /** ピタゴラス数であるか判定 */
    private static boolean isPitagorasNumber(int a, int b, int c) {
        return c * c == a * a + b * b;
    }
}

補足

num() メソッドは無限に繰り返される数を生成する。(num().forEach,,,　などとすると無限ループに陥る。)
num() に中間操作を適用して作成された generate() メソッドの戻り値も無限であることに注意。

このようにStream APIを使うことで、本来無限であるはずのピタゴラス数を自然に表現できる。
呼び出し側では、件数制限の終端操作を適用する必要がある。
これを手続き型のコードで記述しようとすると、無限ループの中で条件が一致した場合にbreakするような判定処理にしなければならない。

[前多 賢太郎]