ScalaでProject Euler（116）

分割数p(k)には、実は次のような漸化式があります。

p(k) = p(k - 1) + p(k - 2) - p(k - 5) - p(k - 7) + p(k - 12) + p(k - 15) - …
p(0) = 1
p(k) = 0 (k < 0)

漸化式のpの中のk - mのmは五角数になっています。符号は、+ + - - + + - -となります。
これで大きな数の分割数についても求めることができます。

import scala.collection.mutable.Map
import scala.testing.Benchmark

val D = 1e6.toInt

def solve() = {
    val ns = Stream.from(1).flatMap(n => Stream(n, -n))
    val polygons = ns.map(n => n * (n * 3 - 1) / 2)
    val signs = Stream.from(0).flatMap(n => Stream(1, 1, -1, -1))
    
    val ps = Map[Int,Int]()
    ps(0) = 1
    
    def calc_p(n :Int) :Int = {
        // p(6) = -p(1) + p(4) + p(5)
        // n : 6 => a : List((1, -1), (4, 1), (5, 1))
        val a = polygons.takeWhile(n >=).map(n -).zip(signs)
        val v = a.map(x => ps(x._1) * x._2).sum % D
        ps(n) = v
        v
    }
    
    println (Iterator.from(1).filter(n => calc_p(n) % D == 0).next)
}

object Test extends Benchmark {
    def run() = solve
}

println (Test.runBenchmark(1))