patternMinor

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

Submitted by: @import:stackexchange-codereview·Mar 10, 2026·

Viewed 0 times

thenumbervectorsinternallyusescalculatingpartitionsstatefulcountinteger

Problem

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  do 
      let i = ((fromIntegral i') :: Integer)
      !sR  ST s ()
          loop k = do
              let f = (fint (i - k * (3 * k - 1) `div` 2))
              when (not (f  s + vec_f) 
                      else do vec_f  s - vec_f) 
                   let f = (fint (i - k * (3 * k + 1) `div` 2))
                   let xx = f :: Int
                   when (not (f  s + vec_f ) 
                          else do vec_f  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

Solution

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 

- 
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 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.

Code Snippets

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

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

Context

StackExchange Code Review Q#147178, answer score: 6

Revisions (0)

No revisions yet.