Source code in the scala programming language

object Main {

  // This is a special class for memoization
  case class Memo[A,B](f: A => B) extends (A => B) {
	  private val cache = Map.empty[A, B]
	  def apply(x: A) = cache getOrElseUpdate (x, f(x))
  }

  // Naive, but memoized solution
  lazy val namesStartingMemo : Memo[Tuple2[Int, Int], BigInt] = Memo {
    case (1, 1) => 1
    case (a, n) => 
	    if (a > n/2) namesStartingMemo(a - 1, n - 1)
	    else if (n < a) 0
	    else if (n == a) 1
	    else (1 to a).map(i => namesStartingMemo(i, n - a)).sum
    
  }
  
  def partitions(n: Int) = (1 to n).map(namesStartingMemo(_, n)).sum
  
  // main method
  def main(args: Array[String]): Unit = {
    for (i <- 1 to 25) {
    	for (j <- 1 to i) { 
	      print(namesStartingMemo(j, i));
	      print(' ');
	    }
    	println()
    }
    println(partitions(23))
    println(partitions(123))
    println(partitions(1234))
    println(partitions(12345))
  }
}


val cache = new Array[BigInt](15000)
cache(0) = 1
val cacheNaive = scala.collection.mutable.Map[Tuple2[Int, Int], BigInt]()

def p(n: Int, k: Int): BigInt = cacheNaive.getOrElseUpdate((n, k), (n, k) match {
    case (n, 1) => 1
    case (n, k) if n < k => 0
    case (n, k) if n == k => 1
    case (n, k) =>
        if (k > n/2) p(n - 1, k - 1)
        else p(n - 1, k - 1) + p(n - k, k)
})

def partitions(n: Int) = (1 to n).map(p(n, _)).sum

def updateCache(n: Int, d: Int, k: Int) =
    if ((k & 1) == 1) cache(n) = cache(n) + cache(d)
    else cache(n) = cache(n) - cache(d)

def quickPartitions(n: Int): BigInt = {
    cache(n) = 0
    for (k <- 1 to n) {
        val d = n - k * (3 * k - 1) / 2
        if (d >= 0) {
            updateCache(n, d, k)

            val e = d - k
            if (e >= 0) {
                updateCache(n, e, k)
            }
        }
    }
    cache(n)
}

for (i <- 1 to 23) {
    for (j <- 1 to i) {
        print(f"${p(i, j)}%4d")
    }
    println
}
println(partitions(23))

for (i <- 1 until cache.length) {
    quickPartitions(i)
}
println(quickPartitions(123))
println(quickPartitions(1234))
println(quickPartitions(12345))