Scala で無限個の無限リストをソートされた形で結合する

以前の記事『Scala で無限リストをソートされた形で結合する』の続き。　前の記事では元記事の本題の「無限個の無限リスト」まで行き着かなかったので、この記事で見ていきます。

Haskell で書かれた元記事はこちら：
itchyny.hatenablog.com
今回扱う問題は5番目、6番目の以下の問題（そう言えば、前回の記事では問題を引用し忘れてました）：

正の整数のうち、ある奇素数pに対するp角数の要素であるような数の集合を考える。このような数の中で、小さい方から1万番目の数字を答えよ。

21は三角数であり八角数であり21角数であるため、多角数として三通りの表現ができる。多角数としてちょうど五通りの表現ができる数のうち、小さい方から150番目の数字を答えよ。

素数列

まずは素数を要素とする素数列を定義しましょう。　前回は整数を Long で定義していましたが、今回扱うサンプルコードでは Int で充分なので整数として Int を使います。　この辺りから、拙者の Haskell 力では元記事の Haskell コードを正確に読み取るのが困難になってきたので、微妙に異なるアルゴリズムになってるところがあるかもしれません。

さて、素数列に戻りましょう。元記事によると Haskell だと素数列を返すコードは2行で書けるとのことなので、Scala でも頑張ってみましょう：

  // 素数列
  val primes: Stream[Int] = 2 #:: Stream.from(3, 2).filter(isPrime)
  def isPrime(n: Int): Boolean = n > 1 && primes.takeWhile(p => p*p <= n).forall(n % _ != 0)

Stream.from(3, 2) は3から始めて2ずつ増える等差数列、まぁ3以上の奇数列ですね。
isPrime は素数判定をするメソッドで、引数の平方根以下の素数全てで割り切れない場合に真を返します。

一応、Scala でも2行で書けました！　実行してみると

  assert( primes.take(10) == Seq(2, 3, 5, 7, 11, 13, 17, 19, 23, 29) )

10個程度だと不安ですが、まぁいいでしょう。

多角数列

次は多角数を要素とする数列。　 ${ n }$ 角数は初項1、公差 ${ n-2 }$ の等差数列を足し上げたものになっています。　奇数の和が平方数になっているのの拡張ですね。　Scala では、遅延評価の数列の足し上げは Stream.scanLeft メソッドでできます：

  // 多角数
  def polygonals(k: Int): Stream[Int] = Stream.from(k-1, k-2).scanLeft(1)(_ + _)

Scala では、Haskell にある初期値を与えなくていい scan1 にあたるものがないようなので、scanLeft に初期値1を与えつつ、足し上げる数列（Stream.from で生成される）は初項を k-1 （本来は第2項）からにしています。

実際に動かしてみると

  assert( polygonals(3).take(10) == Seq(1, 3, 6, 10, 15, 21, 28, 36, 45, 55) )
  assert( polygonals(4).take(10) == Seq(1, 4, 9, 16, 25, 36, 49, 64, 81, 100) )

となって大丈夫そうです。

無限個の無限リストの和集合

では、本題の「無限個の無限リストの和集合」を生成するメソッドを見ていきましょう。　「無限個の無限リスト」と言っても何でもいいわけではなく、各無限リストは単調に増加し*1、また、各数列族の同じ番号の項を比べてもやはり単調に増加している場合に限ります。

例えば、多角数列族（ ${ n }$ 角数の ${ n }$ を族パラメータとして）を見たとき、各多角数列は順々に大きくなる単調増加数列で、また、多角数列の特定の番号の項（たとえば第5項）を順に見ていったときも、族番号が増えるに従って増加していくので、多角数列の族は条件を満たします。　元記事では表で説明してあるのでそちらの方が分かりやすいかと思います。

アルゴリズムの説明も元記事の方で表を使って説明されているので、ここでは省略。

さて、元記事の Haskell コードを拙い Haskell 読解力で Scala コードに書き換えて見たのですが、元記事で定義されている extract 関数が末尾再帰になっていないせいで、Scala では StackOverflow を起こしてしまいます（小さい番号のところでは動いてたのでコードは間違ってないと思いますが）。　また、extract にあたる関数を末尾再帰で書き換えてみたのですが、（問題にある）10000番目の項を取得しようとすると、実行環境のせいか処理が返ってこないので、少々コードを変更しまています。

元記事での union' 関数は、（Scala では ' を名前として使えないので）ここでは multiunion メソッド*2としましょう。　実装はこんな感じ：

  def multiunion[A: Ordering](xss: Stream[Stream[A]]): Stream[A] = {

    // k : k個の Stream を取り出して次の要素を探す
    def go(k: Int, xss: Stream[Stream[A]]): Stream[A] = {
      val (newK, m, yss) = extractMin(k, xss)
      m #:: go(newK, yss)
    }

    // 返り値 : (次の k, head の最小値, 最小値を除去した数列族)
    def extractMin(k: Int, xss: Stream[Stream[A]]): (Int, A, Stream[Stream[A]]) = {
      val (m, i) = xss.take(k).map(_.head).zipWithIndex.minBy(_._1)
        // xss の最初の k 個の Stream に対して
        // 最小値となる head の値とその Stream のインデックスを取得

      (k max (i+2), m, modifyStream(i, xss))
    }

    // i : head を削除する Stream のインデックス
    def modifyStream(i: Int, xss: Stream[Stream[A]]): Stream[Stream[A]] = {
      val (first, second) = xss.splitAt(i)
      first ++: second.head.tail +: second.tail
    }

    go(1, xss)
  }

go メソッドと extractMin メソッドの動作は、元記事の同名関数と同じです。
元記事の extract 関数を末尾再帰に書き換える代わりに、ここでは最小値がどの Stream の head であるか探すのは最小値の計算と一緒に行い、数列族（Stream[Stream[A]] 型のオブジェクト）からこの最小値（結果の数列の次の要素）を削除する処理を modifyStream メソッドにしています*3。
minBy メソッドは引数の変換を施した結果が最小になる最初の要素を返します。
最小値のインデックスに使っている i は、元記事では l+1 にあたるものです。

無限個の無限リストと言っても、（上記の条件を満たすなら）上手くやれば次の要素の候補となるのは有限個だよ、というのがミソなのですね。

1つ目の問題

無限個の無限リストの和集合を返すメソッドができたので、記事最初に書いた問題の解決に移りましょう。　まずは1つ目（元記事の5つ目）の問題：

正の整数のうち、ある奇素数pに対するp角数の要素であるような数の集合を考える。このような数の中で、小さい方から1万番目の数字を答えよ。

まずは multiunion メソッドを使って、奇素数の多角数の和集合を取得してみましょう：

  val seq0 = multiunion(primes.tail.map(polygonals(_).tail))
  assert( seq0.take(30) == 
    Seq(3,5,6,7,10,11,12,13,15,17,18,19,21,22,23,28,29,30,31,34,35,36,36,37,41,43,45,47,48,51) )

primes は素数列ですが、最初の項は2で、今必要なのは奇素数なので、primes.tail としています。
polygonals の初項は全て1で、これは無限個あってプログラムが終わらないので、polygonals(_).tail で1を除いています。

実行結果は1を除いて奇素数の多角数の和集合になっています。　ただし、36のように複数の多角数で現れる（36の場合、三角数の第8項と十三角数の第3項）ものはその回数だけ重複して現れています（なので multiunion と名付けました）。　この重複を除くのは、Haskell では簡単なようですが、Scala では同じようなメソッド distinct を使うと無限 Stream の場合処理が返ってこないので、仕方なく手で書きます。　メソッド名は group で：

  def group[E](seq: Stream[E]): Stream[(E, Int)] = {
    val head = seq.head
    val grouped = seq.takeWhile(_ == head)
    (head, grouped.size) #:: group(seq.drop(grouped.size))
  }

返される Stream の要素は、元の Stream の要素とその重複数のタプルです（元記事の Haskell のコードとはちょっと違いますが）。　先ほどのサンプルコードの seq0 に group を作用させると以下のようになります：

  val seq1 = group(seq0)
  assert( seq1.slice(20, 30) ==
    Seq((35,1),(36,2),(37,1),(41,1),(43,1),(45,1),(47,1),(48,1),(51,1),(53,1)) )

きちんと36の重複数が2になってますね。　重複を落とした Stream を得たい場合は、この Stream の要素になっているタプルの第1要素をとればいいだけです。　よって、multiunion メソッドと group メソッドを併せて、重複のない和集合を返すメソッド union は以下のように定義できます：

  def union[A: Ordering](xss: Stream[Stream[A]]): Stream[A] = group(multiunion(xss)).map(_._1)

これを使って、奇素数の正多角数の和集合は（最初に除いていた1を付け足して）以下で得られます：

  val primePolygonals = 1 #:: union(primes.tail.map(polygonals(_).tail))

実際に動かしてみると

  assert( primePolygonals.take(30) ==
    Seq(1,3,5,6,7,10,11,12,13,15,17,18,19,21,22,23,28,29,30,31,34,35,36,37,41,43,45,47,48,51) )

きちんと36の重複がなくなっていますね。　では、この primePolygonals を使って、問題にある奇素数の正多角数の10000番目の要素を計算してみましょう。　拙者の環境では処理に結構時間がかかりましたが

  assert( primePolygonals(10000-1) == 40641 )

となって元記事の結果と一致しています。

2つ目の問題

では2つ目の問題（元記事では6つ目）もいってみましょう：

21は三角数であり八角数であり21角数であるため、多角数として三通りの表現ができる。多角数としてちょうど五通りの表現ができる数のうち、小さい方から150番目の数字を答えよ。

今度は奇素数ではなく全ての多角数（3以上）が対象です。　重複の個数を探すので、union メソッドではなく multiunion を group した結果の（タプルの）Stream を使うとコードは簡単に書けます。　まずは特定の重複の個数の要素を返す polygonalsIntersection(Int) メソッド。　探したい重複の個数を引数として指定します：

  def polygonalsIntersection(n: Int): Stream[Int] =
    group(multiunion(Stream.from(3).map(polygonals(_).tail))).filter(_._2 == n).map(_._1)

ちょっとゴチャゴチャしてますが、group(multiunion(...)) で正多角数族の和集合（で各要素が重複数とタプルにされている）を生成し、重複数が n となっているものをフィルタリングしています。　最初数項を試してみると

  assert( polygonalsIntersection(5).take(6) == Seq(225, 231, 276, 325, 435, 441) )

となって、元記事と一致しています。　問題となっている150番目の数を求めると、1問目以上に時間がかかりますが

  assert( polygonalsIntersection(5)(150-1) == 10206 )

となって、合ってそうです。　polygonalsIntersection の各要素について、重複がどの多角数由来なのかを見るコードも書いておきましょう。

  // 単調に増加する無限 Stream に対して
  // 要素を含んでいるかどうかを検査するメソッド
  def elem(x: Int, xs: Stream[Int]): Boolean =
    x == xs.dropWhile(_ < x).head

  val n = 5
  for(p <- polygonalsIntersection(n).take(10)){
    for(q <- 3 to p; if elem(p, polygonals(q)))
      yield (p, q)
  } grouped n map (_.mkString("[", ", ", "]")) foreach println

ちょっと for 文の使い方が妙な気もしますが、とりあえずそれなりの出力が得られればよしとしましょう。　実行結果は

[(225,4), (225,8), (225,24), (225,76), (225,225)]
[(231,3), (231,6), (231,17), (231,78), (231,231)]
[(276,3), (276,6), (276,20), (276,93), (276,276)]
[(325,3), (325,6), (325,9), (325,34), (325,325)]
[(435,3), (435,6), (435,45), (435,146), (435,435)]
[(441,4), (441,14), (441,31), (441,148), (441,441)]
[(540,7), (540,10), (540,21), (540,181), (540,540)]
[(595,3), (595,15), (595,30), (595,61), (595,595)]
[(616,7), (616,13), (616,31), (616,104), (616,616)]
[(651,5), (651,9), (651,45), (651,218), (651,651)]

となって、大丈夫そうですね。

まとめ

Scala では Haskell ほど再帰の最適化をしてくれないので少々アルゴリズムを変更しないといけないところがありましたが、まぁ、うまくすれば無限個の無限リスト (Stream) も扱えることが分かりました。　元記事を見ると Haskell のコードは結構数学っぽく書けるようなので、数学好きの人には好まれそうですね。　Scala で書き換えると、全体的にコード量は多くなってるようです（Scala のプロが書けばもっと短くなるかもしれませんが）。

あと、コードを書くときの思考の順序が逆になってるところがいくつもあったのがちょっと面白かったです。　これは Haskell と Scala の文法上の違いのせいなのか、拙者の脳がいまいち関数型プログラミングになっていないせいなのか分かりませんが、Haskell のコードがいまいちスッと読めないのはこの辺りに原因がありそう。