ScalaでProject Euler（23）

Problem 12

引き続き約数の個数がある数より大きい自然数を昇順に列挙する（1）でPythonで書いたコードをScalaにします。

Option

指数の組合せに1加えて次の指数の組合せするPythonのコードは次のようでした。

for k in xrange(1, len(es0)):
    if es0[k-1] == es0[k] + 1:
        es2 = es0[:k] + (es0[k] + 1,) + es0[k+1:]
        pq.put((value(es2), es2))
        break
    elif es0[k-1] > es0[k]:
        break
else:
    if es0[-1] == 1:
        es2 = es0 + (1,)
        pq.put((value(es2), es2))

自分で書いておいてなんですが、わかりにくいコードですね。人間デバッガになってステップ実行しないとわからないですね。
Scalaでは関数型らしくHaskellっぽく書きましょう。再帰にしてOption型を使います。Option型はHaskellのMaybeモナドみたいなものです。

def next_combination(es :Option[List[Int]]) :Option[List[Int]] = es match {
    case Some(1 :: tail) => Some(1 :: 1 :: tail)
    case Some(e1 :: e2 :: tail) if e1 == e2 + 1 => Some(e1 :: e1 :: tail)
    case Some(e1 :: e2 :: tail) if e1 == e2 =>
        next_combination(Some(e2 :: tail)) match {
            case Some(b) => Some(e2 :: b)
            case None => None
        }
    case _ => None
}

順番だけちゃんと考えて、あとは書くだけです。すっきりしています。

println (next_combination(Some(List(3, 2, 1))))     // Some(List(3, 3, 1))
println (next_combination(Some(List(2, 2, 1))))     // Some(List(2, 2, 2))
println (next_combination(Some(List(1, 1, 1))))     // Some(List(1, 1, 1, 1))
println (next_combination(Some(List(3, 1, 1))))     // None

一発でちゃんと動きました。
最後の順列に0をはさむのも同様に書けます。ただし、これはListの右のほうが重要なのでreverseして処理してreverseします。

def next_value(es :List[Int]) = {
    def f(es :Option[List[Int]]) :Option[List[Int]] = es match {
        case Some(0 :: 0 :: tail) => None
        case Some(0 :: e :: tail) => Some(e :: 0 :: tail)
        case Some(e :: tail) => f(Some(tail)) match {
            case Some(b) => Some(e :: b)
            case None => None
        }
        case _ => None
    }
    
    f(Some(es.reverse)) match {
        case Some(a) => Some(a.reverse)
        case None => None
    }
}

println (next_value(List(3, 2, 1)))         // None
println (next_value(List(3, 0, 2, 1)))      // Some(List(0, 3, 2, 1))
println (next_value(List(3, 2, 0, 0, 1)))   // None