倭マン's BLOG

くだらない日々の日記書いてます。 たまにプログラミング関連の記事書いてます。 書いてます。

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

あとで読もうと思って数ヶ月放置してた記事
itchyny.hatenablog.com
を消化しようと、途中まで Scala で書いてみた。 まぁ、ホント書き換えただけの記事。 しかも、本論の「無限個の無限リスト」まではまだ行けず、その手前の3個の無限リストまで。

ところで、Haskell ってコードがかなり数学っぽく書けるんだなぁ。

記事の内容

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

まずは三角数と平方数の定義から。 後のサンプルでいくつか Int の範囲を超えるので、整数は Long にしておきます*1

  val integers = Stream from 1 map (_.toLong)
  assert( integers.take(5) == Seq(1, 2, 3, 4, 5) )

  // 三角数
  val triangulars = integers map (n => n * (n + 1) / 2)
  assert( triangulars.take(5) == Seq(1, 3, 6, 10, 15) )

  // 平方数
  val squares = integers map (n => n * n)
  assert( squares.take(5) == Seq(1, 4, 9, 16, 25) )
  • Stream.from(Int) は引数の整数から始めて、1ずつ増える整数の(無限)Stream を生成するメソッドですね。

さて、これを踏まえて、2つの無限リストの和集合を返すメソッド union2 を定義してみましょう:

  def union2(xs: Seq[Long], ys: Seq[Long]): Stream[Long] = (xs, ys) match {
    case (_, Nil) => xs.toStream  // 一方が空なら
    case (Nil, _) => ys.toStream  // もう一方を返す
    case _ =>
      // 2つの無限リストの和集合を作る
      val (x, y) = (xs.head, ys.head)
      x compare y match {
        case n if n < 0 => x #:: union2(xs.tail, ys)
        case 0          => x #:: union2(xs.tail, ys.tail)
        case _          => y #:: union2(xs, ys.tail)
      }
  }
  • 最初の方の case 文は、引数の Seq のどちらか一方が空なら他方を返すためのコード。
  • 2つの無限リストの和集合を作る部分では、2つのリストの先頭の要素を比べて、小さい方を先頭にした新たな Stream を作っています。 再帰的な定義の Stream を遅延評価できちんと動くようにするために Stream の「#::」演算子(メソッド)を使っています。 この演算子を使いたいので、返り値の型は Seq ではなく Stream にしてあります。

Haskell のコードに比べてちょっと長くなって、メソッドの宣言などとも相まってごちゃごちゃしてる感がありますが、まぁ仕方ないでしょう。

では動作確認。 まずは一方が空の場合:

  assert( union2(triangulars, Nil)(1000-1) == 500500 )  // 三角数の1000番目
  assert( union2(Nil, squares)(1000-1) == 1000000 )  // 平方数の1000番目

「1000-1」は Stream の番号が0から始まるところからくるズレです。 まぁ、動作に問題はなさそうですね。 テストとしては両方空集合の場合とか、一方が有限集合の場合とか要りそうですが、ここではよしとしましょう。 次は2つの無限リストの場合:

  assert( union2(triangulars, squares)(1000-1) == 173166 )
  assert( union2(triangulars, squares)(100000-1) == 1716013236 )
  assert( union2(triangulars, squares)(1000000-1) == 171575840736L)

となって、うまく動いてます。 処理時間も特に気になりません。

2個の無限リストの共通部分

和集合の次は共通部分。 メソッド名は intersect2 としましょう。 基本的な構造は和集合の場合と同じです:

  def intersect2(xs: Seq[Long], ys: Seq[Long]): Stream[Long] = (xs, ys) match {
    case (_, Nil) => Stream.empty  // 一方が空なら
    case (Nil, _) => Stream.empty  // 空の Stream を返す
    case _ =>
      // 2つの無限集合の共通部分を作る
      val (x, y) = (xs.head, ys.head)
      x compare y match {
        case n if n < 0 => intersect2(xs.tail, ys)
        case 0          => x #:: intersect2(xs.tail, ys.tail)  // 共通要素を先頭要素に
        case _          => intersect2(xs, ys.tail)
    }
  }
  • 定義的に当たり前ですが、2つの引数の先頭要素が同じ時だけ #:: 演算子で先頭要素を付け加えています。 それ以外の場合は、その要素を落として再度 intersect2 メソッドを呼び出しています。

特に目新しいものはないですね。 では動作確認:

  assert( intersect2(triangulars, squares)(6-1) == 48024900 )
  assert( intersect2(triangulars, squares)(7-1) == 1631432881 )

よしよし。

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

次は3個の無限リストの場合。 まずはサンプルに使う五角数の定義から。

  // 五角数
  val pentagonals = integers map (n => n * (3 * n - 1) / 2)
  assert( pentagonals.take(5) == Seq(1, 5, 12, 22, 35) )

よく考えると、三角数や平方数の定義って図形的(幾何学的)なところからきてるんだろうけど、五角数って敷き詰められないんじゃ?という疑問が湧いてくるけど、Wikipedia によると実際あんまり綺麗な並べ方ではないようで。

まぁそれはいいとして、3個の無限リストの和集合を作るメソッド union3 の実装はこんな感じ:

  def union3(xs: Stream[Long], ys: Stream[Long], zs: Stream[Long]): Stream[Long] = {
    val w = Seq(xs.head, ys.head, zs.head).min
    w #:: union3(xs.dropWhile(_ == w), ys.dropWhile(_ == w), zs.dropWhile(_ == w))
  }

元記事でやっているように、引数が無限リストであることを前提に、引数の空チェックを省いてます。 それを考慮しても、なんか2個の場合よりかなり短くなってますねぇ。 やっていることは

  • 引数の無限リストの先頭要素のうち、最小のもの(これが次の要素となる)を求める => w
  • w を先頭にし、各無限リストから w を除いたものを引数に union3 を再帰的に呼び出し tail にする

という感じです。

では動作確認:

  assert( union3(triangulars, squares, pentagonals)(10000-1) == 9603153 )
  assert( union3(triangulars, squares, pentagonals)(100000-1) == 958335849 )

この union3 は union2 と違って、引数に Nil とか Seq を渡すとコンパイルエラーになりますが、有限リストを Stream にして渡しても実行事例外がでるだけでいいことはありません。

3個の無限リストの共通部分

今回の最後は3個の無限リストの共通部分。 まずはサンプルで使う六角数*2の定義から:

  val hexagonals = integers map( n => n * (2 * n - 1))
  assert( hexagonals.take(5) == Seq(1, 6, 15, 28, 45) )

これまたどうでもいいことだけど、六角数を図形的に見たとき、穴がないように敷き詰めるわけではないんですなぁ。 今まで六角数とか使ったことなかったので特に問題はないんだけど、誤解を招きそうな定義や。 七角数以上は平坦な平面上では描けないっすねぇ。

では3個の無限リストの共通部分を作るメソッド intersect3 を実装しましょう。 ここでもやはり引数が無限リストであることを前提としてます。

  def intersect3(xs: Stream[Long], ys: Stream[Long], zs: Stream[Long]): Stream[Long] = {
    val (x, y, z) = (xs.head, ys.head, zs.head)
    () match {  // 行頭の Unit () は特に意味ナシ
      case _ if x < y || x < z => intersect3(xs.tail, ys, zs)
      case _ if y < z || y < x => intersect3(xs, ys.tail, zs)
      case _ if z < x || z < y => intersect3(xs, ys, zs.tail)
      case _ => x #:: intersect3(xs.tail, ys.tail, zs.tail)  // x, y, z が同じ場合
    }
  }
  • 条件分岐を match-case 文で書きたかったので Unit オブジェクト () に対して match を行ってますが、別に他のオブジェクトでもかまいません。
  • 3個の引数の先頭要素 x, y, z が同じときだけ先頭要素を付け加えています。

動作確認コードはこんな感じ:

  assert( intersect3(triangulars, squares, hexagonals)(3-1) == 1413721 )

特に問題なく動いてます。

元記事の本題は、タイトルにあるように無限個の無限リストを扱う話ですが、この記事では有限個の無限リストでいっぱいいっぱいな感じ。 続き、というか本題もそのうちやるつもりではありますが、いつになることやら。

【修正】
2個の無限リストの和集合と共通部分のコードで、compare 結果が -1 となる case を負の case となるように修正しました(-1でも一応動きますが)。
すごいHaskellたのしく学ぼう!

すごいHaskellたのしく学ぼう!

*1:とりあえず今回のサンプルコードでは Long で事足りますが、後続記事では BigInt などにするかもしれません。 型クラス使ってジェネリックなコードにすることもできますが、まぁ今回はやめときます。

*2:三角数、平方数、五角数の共通部分を求めようとすると2番目の要素を求めるだけでも時間がかかる(というか存在するのかどうか知らないけど)ので。