ScalaでProject Euler（80）

Problem 50

PriorityQueueを使えば簡単です。まず、2からはじまる最長の素数列を求めます。そして、(列の長さ, 最初の素数のインデックス, 列の和)というタプルをPriorityQueueに入れます。PriorityQueueから要素を取り出して、今までの最長の素数未満の長さならそこで終わりです。そうでなくそれが素数なら最長を更新します。素数でなければはじまりは同じで一つ短い列をPriorityQueueに入れます。その列がその素数からはじまって最長の列なら、次の素数からはじまる最長の列もPriorityQueueに入れます。
10億まででも一瞬です。

import scala.collection.mutable.ArrayBuffer
import scala.collection.mutable.PriorityQueue

object Primes {
    private val primes = new ArrayBuffer[Int]
    private val L = 1e5.toInt
    private var offs = 0
    
    def apply(n :Int) = {
        if(n >= primes.size)
            sieve
        primes(n)
    }
    
    private def sieve() = {
        if(offs == 0)
            sieve_1st
        else
            sieve_2nd_or_later
        offs += L
    }
    
    private def sieve_1st() = {
        val a = new Array[Boolean](L)
        for(i <- 2 to (L - 1)) a(i) = true
        for(p <- 2 to (L - 1) if a(p);
                    m <- (p * 2) to (L - 1) by p)
            a(m) = false
        
        primes ++= (2 to (L - 1)).filter(a(_)).toArray
    }
    
    private def sieve_2nd_or_later() = {
        val N = offs + L - 1
        val a = new Array[Boolean](L)
        for(i <- 0 to (L - 1)) a(i) = true
        for(p <- primes.takeWhile(p => p * p <= N);
                    m <- ((offs + p - 1) / p * p) to N by p)
            a(m - offs) = false
        
        primes ++= (offs to N).filter(k => a(k - offs)).toArray
    }
}

def is_prime(n :Int) =
    Iterator.from(0).map(k => Primes(k)).takeWhile(p => p * p <= n).
                                                forall(p => n % p != 0)

type Queue = (Int,Int,Int)  // (列の長さ, 最初の素数のインデックス, 列の和)

// (長さ, 和, 素数か)
class gen_queue extends Iterator[(Int,Int,Boolean)] {
    val o = Ordering[Long].on[Queue](_._1)
    val pq = new PriorityQueue[Queue]()(o)
    init
    
    def hasNext() = true
    
    def next() = {
        val (l0, k0, s0) = pq.dequeue
        val b = is_prime(s0)
        if(!b && l0 > k0)
            pq.enqueue((l0 - 1, k0, s0 - Primes(l0 - k0 - 1)))
        
        if(s0 + Primes(k0 + l0) >= N) { // その素数から始まって最長の列
            val s1 = s0 - Primes(k0) + Primes(k0 + l0)  // 同じ長さの和
            if(s1 < N)          // 2, 3, 5, 7 -> 3, 5, 7, 11
                pq.enqueue((l0, k0 + 1, s1))
            else                // 2, 3, 5, 7 -> 3, 5, 7
                pq.enqueue((l0 - 1, k0 + 1, s0 - Primes(k0)))
        }
        (l0, s0, b)
    }
    
    def init() = {
        // (和, 最大のインデックス)
        val (s0, k0) = Iterator.iterate((0, -1))(s =>
                    (s._1 + Primes(s._2 + 1), s._2 + 1)).
                    dropWhile(_._1 < N).next
        pq.enqueue((k0, 0, s0 - Primes(k0)))
    }
}

val N = 1e6.toInt
val it = new gen_queue
println (Iterator.iterate((0, 0)){x =>
    val (l0, s0) = x
    val (l, s, b) = it.next
    if(b)
        if(l == l0) (l, s.max(s0)) else if(l > l0) (l, s) else (-1, s0)
    else
        if(l >= l0) (l0, s0) else (-1, s0)
}.dropWhile(x => x._1 != -1).map(x => x._2).next)