Scala implementation of @rolfl's integer-partitioning algorithm

Question

While reviewing @Martijn's Distinct sums of integers (finding all combinations of positive integers that sum to a given number), I proposed the following solution:

/** List of non-increasing positive integers */
type Sum = List[Int]

def sums(n: Int): Iterator[Sum] = {
  /** Prepends the head of the stack repeatedly until a total amount has
    * been added.  The leading element of the result is then increased
    * so that the total added is exactly the specified amount.
    */
  @tailrec
  def topUp(amount: Int, stack: List[Int]): List[Int] = {
    if (amount < stack.head)
      (stack.head + amount) :: stack.tail
    else
      topUp(amount - stack.head, stack.head :: stack)
  }

  /** Returns the next partition by popping the first element, increasing the
    * next element, and topping it up to the same total.  Returns None if
    * the input has only one element.
    */
  def advance(partition: Option[Sum]): Option[Sum] = {
    partition.get match {
      case List(n) => None
      case p       => Some(topUp(p.head - 1, p.tail.head + 1 :: p.tail.tail))
    }
  }

  Iterator.iterate(Some(List.fill(n)(1)): Option[Sum])(advance)
          .takeWhile(! _.isEmpty)
          .map(_.get)
}

Example run:

scala> sums(5).foreach(println)
List(1, 1, 1, 1, 1)
List(2, 1, 1, 1)
List(3, 1, 1)
List(2, 2, 1)
List(4, 1)
List(3, 2)
List(5)

I'm pretty satisfied with the algorithm, as it is based on a Java solution that was devised by @rolfl.

However, I'm not happy about the way I've expressed it in Scala, particularly with regards to the termination of the iterator.

• Is there a way to write topUp without explicit recursion, yet still avoiding ::: concatenation?
• Is the : Option[Sum] type annotation on Some(List.fill(n)(1)) avoidable?
• There seems to be a lot of Option / None / Some noise in advance. Is there a better way to handle the termination condition? (Perhaps by using lift?)
• Similarly, Iterator.iterate(…)(…).takeWhile(! _.isEmpty).map(_.get) seems cumbersome. Is there a less verbose way to say "repeatedly apply this function until exhausted" using the standard Scala library?

---

For the record, instead of advance and Iterator.iterate(…).takeWhile(…).map(…), I had previously written it this way, using explicit recursion:

def makeSums(partition: Sum): Stream[Sum] = {
  partition #:: (partition match {
    case List(n) => Stream.Empty
    case p       => makeSums(topUp(p.head - 1, p.tail.head + 1 :: p.tail.tail))
  })
}

makeSums(List.fill(n)(1))

… which returned a Stream[Sum] instead of an Iterator[Sum]. Note that it's not tail-recursive. (I also tried to write it as a self-referential stream, but couldn't figure out how to terminate it.)

I think that, despite being longer, the Iterator.iterate(…) solution is better because it decouples the "repeated application" aspect from the "produce the next result" aspect. Your opinion on this matter would be appreciated.

Answer 1

_{Keep in mind that I likely know less about scala than you and that this is my first review ;)}

Is there a way to write topUp without explicit recursion, yet still avoiding ::: concatenation?

I don't think so, but I may be wrong.

Is the : Option[Sum] type annotation on Some(List.fill(n)(1)) avoidable?

Yes! You can use the Option type constructor like this:

Iterator.iterate(Option(List.fill(n)(1)))(advance)

There seems to be a lot of Option / None / Some noise in advance. Is there a better way to handle the termination condition? (Perhaps by using lift?)

I think so, but not with lift. This just uses the standard map and filter (please let me know if this is incorrect logic).

def advance(partition: Option[Sum]): Option[Sum] =
   partition.filter(_.size > 1)
            .map(p => topUp(p.head - 1, p.tail.head + 1 :: p.tail.tail))

Similarly, Iterator.iterate(…)(…).takeWhile(! _.isEmpty).map(_.get) seems cumbersome. Is there a less verbose way to say "repeatedly apply this function until exhausted" using the standard Scala library?

I really don't know.

One minor point. Here

Iterator.iterate(Some(List.fill(n)(1)): Option[Sum])(advance)
        .takeWhile(! _.isEmpty)
        .map(_.get)

You use ! _.isEmpty, but option has a perfect method for you, isDefined.