Skip to content

Instantly share code, notes, and snippets.

@phlummox

phlummox/part.hs

Last active Nov 16, 2016
Embed
What would you like to do?
calculate count of integer partitions of a number, in Haskell
{-# 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
result
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment