ScalaでProject Euler(24)

Problem 12

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

Set

Setはこんな感じに使います。

import scala.collection.mutable.Set

val s = Set[Int]()
s += 3
s += 1
s += 4
println (s)
println (s.contains(3))

これでProblem 12が書けます。
しかし、遅いですね。

(追記)
少し間違えていただけでした。修正しました。

import scala.collection.mutable.ArrayBuffer
import scala.collection.mutable.PriorityQueue
import scala.collection.mutable.Set
import Math.sqrt

object Prime {
    val primes = new ArrayBuffer[Int]
    primes += 2
    
    def apply(k :Int) =
        if(k < primes.size)
            primes(k)
        else {
            for(n <- Stream.from(primes(primes.size-1) + 1).
                    takeWhile(n => k >= primes.size) if is_prime(n))
                primes += n
            primes(primes.size-1)
        }
    
    def is_prime(n :Int) =
        primes.takeWhile(p => p * p <= n).forall(n % _ != 0)
}

def pow(n :Int, e :Int) =
    (1 to e).foldLeft(1L)((x, y) => x * n)

def value(es :List[Int]) =
    es.zip(0 to (es.size - 1)).foldLeft(1L)(
                (x, y) => x * pow(Prime(y._2), y._1))

def num_divs(es :List[Int]) =
    es.foldLeft(1)((x, y) => x * (y + 1))

class gen_combinations extends Iterator[(Long,List[Int])] {
    val o = Ordering[Long].on[(Long,List[Int])](-_._1)
    val pq = new PriorityQueue[(Long,List[Int])]()(o)
    pq.enqueue((2, List(1)))
    
    def hasNext() = true
    
    def next() = {
        val (v0, es0) = pq.dequeue
        
        val es1 = (es0.head + 1) :: es0.tail
        pq.enqueue((value(es1), es1))
        next_combination(Some(es0)).foreach(es => pq.enqueue((value(es), es)))
        
        (v0, es0)
    }
    
    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)).map(e2 :: _)
        case _ => None
    }
}

def is_descending(es :List[Int]) =
    es.zip(es.tail).forall(x => x._1 >= x._2)

class gen_greater_permutations(N :Int) extends Iterator[(Long,List[Int])] {
    val o = Ordering[Long].on[(Long,List[Int])](-_._1)
    val pq = new PriorityQueue[(Long,List[Int])]()(o)
    val g = (new gen_combinations).filter(x => num_divs(x._2) > N)
    val set_used = Set[Long]()
    
    put(g.next)
    
    def hasNext() = true
    
    def next() = {
        val (v0, es0) = pq.dequeue
        
        val ess = next_permutation(es0).
                    filter(es => !set_used.contains(value(es)))
        for(es <- ess) put((value(es), es))
        
        if(is_descending(es0))
            pq.enqueue(g.next)
        
        (v0, es0)
    }
    
    private def put(x :(Long,List[Int])) = {
        pq.enqueue(x)
        set_used += x._1
    }
    
    def next_permutation(es :List[Int]) :List[List[Int]] = es match {
        case (e1 :: e2 :: tail) if e1 > e2 =>
                (e2 :: e1 :: tail) :: next_permutation(e2 :: tail).map(e1 :: _)
        case (e1 :: e2 :: tail) => next_permutation(e2 :: tail).map(e1 :: _)
        case _ => Nil
    }
}

class gen_greater_numbers(N :Int) extends Iterator[Long] {
    val o = Ordering[Long].on[(Long,List[Int])](-_._1)
    val pq = new PriorityQueue[(Long,List[Int])]()(o)
    val g = new gen_greater_permutations(N)
    
    pq.enqueue(g.next)
    
    def hasNext() = true
    
    def next() = {
        val (v0, es0) = pq.dequeue
        
        val rev_es0 = es0.reverse
        val rev_es1 = rev_es0.head :: 0 :: rev_es0.tail
        val es1 = rev_es1.reverse
        pq.enqueue((value(es1), es1))
        
        if(es0.forall(0 !=))
            pq.enqueue(g.next)
        next_value(es0).foreach(es => pq.enqueue((value(es), es)))
        
        v0
    }
    
    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)).map(e :: _)
            case _ => None
        }
        
        f(Some(es.reverse)) match {
            case Some(a) => Some(a.reverse)
            case None => None
        }
    }
}

def is_triangle(n :Long) = {
    val m = sqrt(n * 2.0).toLong
    n == m * (m + 1) / 2
}

val t0 = System.currentTimeMillis
val N = 50000
val g = new gen_greater_numbers(N)
println (g.filter(is_triangle).next)
println ((System.currentTimeMillis - t0) + "ms")