Integer partitioning in Scala

Question

Given n ( say 3 people ) and s ( say 100$ ), we'd like to partition s among n people.

So we need all possible n-tuples that sum to s

My Scala code below:

def weights(n:Int,s:Int):List[List[Int]] = {
     List.concat( (0 to s).toList.map(List.fill(n)(_)).flatten, (0 to s).toList).
     combinations(n).filter(_.sum==s).map(_.permutations.toList).toList.flatten
}

println(weights(3,100))

This works for small values of n. ( n=1, 2, 3 or 4).

Beyond n=4, it takes a very long time, practically unusable.

I'm looking for ways to rework my code using lazy evaluation/ Stream.

My requirements : Must work for n upto 10.

Warning : The problem gets really big really fast. My results from Matlab -

---For s =100, n = 1 thru 5 results are ---
n=1 :1 combinations
n=2 :101 combinations
n=3 :5151 combinations
n=4 :176851 combinations
n=5: 4598126 combinations
---

Answer 1

You need dynamic programming, or memoization. Same concept, anyway.

Let's say you have to divide s among n. Recursively, that's defined like this:

def permutations(s: Int, n: Int): List[List[Int]] = n match {
  case 0 => Nil
  case 1 => List(List(s))
  case _ => (0 to s).toList flatMap (x => permutations(s - x, n - 1) map (x :: _))
}

Now, this will STILL be slow as hell, but there's a catch here... you don't need to recompute permutations(s, n) for numbers you have already computed. So you can do this instead:

val memoP = collection.mutable.Map.empty[(Int, Int), List[List[Int]]]
def permutations(s: Int, n: Int): List[List[Int]] = {
  def permutationsWithHead(x: Int) = permutations(s - x, n - 1) map (x :: _)

  n match {
    case 0 => Nil
    case 1 => List(List(s))
    case _ => 
      memoP getOrElseUpdate ((s, n), 
                             (0 to s).toList flatMap permutationsWithHead)
  }
}

And this can be even further improved, because it will compute every permutation. You only need to compute every combination, and then permute that without recomputing.

To compute every combination, we can change the code like this:

val memoC = collection.mutable.Map.empty[(Int, Int, Int), List[List[Int]]]
def combinations(s: Int, n: Int, min: Int = 0): List[List[Int]] = {
  def combinationsWithHead(x: Int) = combinations(s - x, n - 1, x) map (x :: _)

  n match {
    case 0 => Nil
    case 1 => List(List(s))
    case _ => 
      memoC getOrElseUpdate ((s, n, min), 
                             (min to s / 2).toList flatMap combinationsWithHead)
  }
}

Running combinations(100, 10) is still slow, given the sheer numbers of combinations alone. The permutations for each combination can be obtained simply calling .permutation on the combination.

Answer 2

Here's a quick and dirty Stream solution:

 def weights(n: Int, s: Int) = (1 until s).foldLeft(Stream(Nil: List[Int])) {
   (a, _) => a.flatMap(c => Stream.range(0, n - c.sum + 1).map(_ :: c))
 }.map(c => (n - c.sum) :: c)

It works for n = 6 in about 15 seconds on my machine:

scala> var x = 0
scala> weights(100, 6).foreach(_ => x += 1)
scala> x
res81: Int = 96560646

As a side note: by the time you get to n = 10, there are 4,263,421,511,271 of these things. That's going to take days just to stream through.

Answer 3

My solution of this problem, it can computer n till 6:

object Partition {
  implicit def i2p(n: Int): Partition = new Partition(n)
  def main(args : Array[String]) : Unit = {
    for(n <- 1 to 6) println(100.partitions(n).size)
  }

}

class Partition(n: Int){
  def partitions(m: Int):Iterator[List[Int]] = new Iterator[List[Int]] {
    val nums = Array.ofDim[Int](m)
    nums(0) = n

    var hasNext = m > 0 && n > 0

    override def next: List[Int] = {
      if(hasNext){
        val result = nums.toList
        var idx = 0
        while(idx < m-1 && nums(idx) == 0) idx = idx + 1
        if(idx == m-1) hasNext = false
        else {
          nums(idx+1) = nums(idx+1) + 1
          nums(0)     = nums(idx) - 1
          if(idx != 0) nums(idx)   = 0
        }
        result
      }
      else Iterator.empty.next
    }
  }
}

1 101 5151 176851 4598126 96560646

---

However , we can just show the number of the possible n-tuples:

val pt: (Int,Int) => BigInt =  {
    val buf = collection.mutable.Map[(Int,Int),BigInt]()
    (s,n) => buf.getOrElseUpdate((s,n),
        if(n == 0 && s > 0) BigInt(0)
        else if(s == 0) BigInt(1)
        else (0 to s).map{k => pt(s-k,n-1)}.sum
        )
  }

  for(n <- 1 to 20) printf("%2d :%s%n",n,pt(100,n).toString)

 1 :1
 2 :101
 3 :5151
 4 :176851
 5 :4598126
 6 :96560646
 7 :1705904746
 8 :26075972546
 9 :352025629371
10 :4263421511271
11 :46897636623981
12 :473239787751081
13 :4416904685676756
14 :38393094575497956
15 :312629484400483356
16 :2396826047070372396
17 :17376988841260199871
18 :119594570260437846171
19 :784008849485092547121
20 :4910371215196105953021