Calculating the count of integer partitions of a number (uses stateful vectors internally)

Question

I've written a Haskell function to count the number of partitions of an integer - it's basically an implementation of the same algorithm as this one, using Euler's Pentagonal Number Theorem:

$$P(n) = \sum_{k=1}^n (-1)^{k+1} \left[ P\left( n - \frac{k\left(3k-1\right)}{2} \right) + P\left( n - \frac{k\left(3k+1\right)}{2} \right) \right]$$

To get it to complete on large-ish inputs (say, 60,000 or so) on a reasonable time , I used a mutable vector, internally. (This cut the time by about half from using an immutable vector, and by a ... very large amount, I forget how much, from using "pure" memoization.) It now calculates \$p(60000)\$ in about 8 seconds or so on my machine, which is reasonable for my purposes. (Though three times slower than a C++ equivalent.)

I wondered if anyone had any suggestions for getting similar performance, but without using mutable state?

Gist

{-# LANGUAGE BangPatterns #-}

import Control.Monad (when, forM_, forM)

import Data.STRef
import Control.Monad.ST

import qualified Data.Vector.Generic.Mutable as GM
import Data.Vector.Generic.Mutable ( write )
import qualified Data.Vector.Mutable as VM


main :: IO ()
main = do
  print $ part 60000

fint :: (Num b, Integral a) => a -> b
fint = fromIntegral

part :: Int -> Integer
part n = runST $ do 
  vec <- VM.replicate (n+1) (-3)
  write vec 0 1
  result <- newSTRef (-1)
  forM_ [1..n] $ \i' -> do 
      let i = ((fromIntegral i') :: Integer)
      !sR <- newSTRef 0 
      let -- loop :: Integer -> ST s ()
          loop k = do
              let f = (fint (i - k * (3 * k - 1) `div` 2))
              when (not (f < 0)) $ do
                   if k `mod` 2 /= 0
                      then do vec_f <- GM.read vec f  
                              modifySTRef' sR (\s -> s + vec_f) 
                      else do vec_f <- GM.read vec f 
                              modifySTRef' sR (\s -> s - vec_f) 
                   let f = (fint (i - k * (3 * k + 1) `div` 2))
                   let xx = f :: Int
                   when (not (f < 0)) $ do
                        if k `mod` 2 /= 0
                          then do vec_f <- GM.read vec f 
                                  modifySTRef' sR (\s -> s + vec_f ) 
                          else do vec_f <- GM.read vec f 
                                  modifySTRef' sR (\s -> s - vec_f) 
                        loop (k + 1)
      loop 1 -- k starts at 1
      !s <- readSTRef sR
      write vec i' s
      when (i' == n) $
        writeSTRef result s
  readSTRef result

Answer 1

You can achieve comparable performance without resorting to mutable state by using IntMap and carefully avoiding unnecessary computations. Let's start with the code below, which is actually the very first Haskell solution at your Pentagonal Number Theorem link:

import Data.IntMap ((!), fromList, insert, findWithDefault)

partition :: Int -> Integer
partition x = if x < 0 then 0 else foldl p (fromList [(0,1)]) [1..x] ! x where
    p s n = insert n (sum [(-1)^(k+1) * (r pred + r succ) | k <- [1..n],
        let r f = findWithDefault 0 (n - div (k * f (3 * k)) 2) s]) s

This code is pretty, concise and expressive, and it's very easy to see that it does implement Euler's formula indeed. Alas, it's not very efficient - on my laptop, it takes more than 5 seconds to compute partition 5000 (I didn't measure how long partition 60000 would take). How can we improve its performance? Well, there's some low-hanging fruit:

• (-1)^(k+1) looks suspect: we don't really need exponentiation actually - what we need is 1 if k is odd and (-1) if k is even.

• k <- [1..n] gives us way too many ks. We should only consider those k values for which at least one of $$n - \frac{k\left(3k-1\right)}{2}$$ and $$n - \frac{k\left(3k+1\right)}{2}$$ is non-negative (because if both of them are negative, then the sum within the square brackets in the formula above is zero). The largest such k is approximately $$\sqrt\frac{2n}{3}$$ (the exact maximum can be computed easily by solving a quadratic equation).

• That is, instead of n possible values for k, we only need to consider fewer than sqrt n values. And that matters a lot because the code performs a lookup (findWithDefault) in the IntMap for each k value considered.

• These two (i.e., the elimination of (-1)^(k+1) and the reduction of the number of lookups) are the biggest wins, but there's some additional minor stuff you can do to improve performance, e.g., you can eliminate one multiplication and one division by noticing that $$n - \frac{k\left(3k+1\right)}{2} = n - \frac{k\left(3k-1\right)}{2} - k$$

• Also, it doesn't hurt to try using a strict IntMap instead of a lazy one.

After implementing all of these suggestions, the code looks like this:

import Data.IntMap.Strict ((!), fromList, insert)

partition :: Int -> Integer
partition x = if x < 0 then 0 else foldl p (fromList [(0,1)]) [1..x] ! x where
    p s n = insert n (sum [sign k * sumOfTwoEarlierPs s n k | k <- [1..maxK n] ]) s
    maxK n = round $ 1/6 + sqrt(1/36 + 2/3 * fromIntegral n)

sign k = if odd k then 1 else (-1)

sumOfTwoEarlierPs s n k = p1 + p2 where
  i1 = n - (k *(k+k+k - 1)) `div` 2
  i2 = i1 - k
  p1 = nonNegativeIndexOnly s i1
  p2 = nonNegativeIndexOnly s i2
  nonNegativeIndexOnly s i = if i >= 0 then s!i else 0

On my laptop, this computes partition 60000 in less than 5 seconds, and its running time is only 5% or so more than that of your code.

Answer 2

You asked for a stable data formulation with comparable speed. My suggestion:

import Data.Array;
part(n)=let{a=zipWith(*)[1,3..][1,4..];c=zipWith(*)[1..][5,11..];
            b=zipWith(*)[1,3..][2,5..];d=zipWith(*)[1..][7,13..];
            y=array(0,n)((0,1):[(i,s(a)+s(b)-s(c)-s(d))|i<-[1..n],
            let s=foldr(+)0.map((y!).(i-)).takeWhile(<=i)])}in(y!n);
main=print(part(60000));

– when compiled with ghc -O to optimize away stack overflows – computes P(60k) in ~12s@800MHz O(n²) runtime practically equal to your expressly stateful computation here!

I agree that an O(1)-accessible array of boxed big integers yields unbeatable minimal cost for sufficient memories when index jumps of up to 200 cells occur in the computation of P(60k).

I hid my computation state behind my index<-[1..n]-sequenced y-recursion though.

And lest I recompute most products of multiplications for array indexing, I memoized them in lazily infinite named sequences

• a := <[1,3,..]*[1,4..]> = (2h+1)(3h+1) = (2h+1)½(6h+3-1) = ½k(3k-1) for odd k=2h+1
• b := <[1,3,..]*[2,5..]> = (2h+1)(3h+2) = (2h+1)½(6h+3+1) = ½k(3k+1) for odd k=2h+1
• c := <[1,2,..]*[5,11..]> = (h+1)(6h-1) = ½(2h+2)(6h-1) = ½k(3k-1) for even k=2h+2
• d := <[1,2,..]*[7,13..]> = (h+1)(6h+1) = ½(2h+2)(6h+1) = ½k(3k+1) for even k=2h+2

and feel rewarded with best code brevity.

P(60k) stack overflows my Debian8 ghc -O of @BenceKodaj's O(log₂n) IntMap structure :(