Haskell, algorithm all possible composition of number

Question

I have a code in haskell which generates three-part composition of number:

kompozycje n = [ (x,y,z) | x<-[1..n], y<-[1..n], z<-[1..n], x+y+z==n]

I would like to make something like kompozycje n k which would generate me k-part compositions and then if for example k would equal 4 there would be four variables and four numbers returned and in condition there would be something like u+x+y+z==n. Is there some simple solution for that?

Answer 1

Yes, yes there is. It uses the list monad and replicateM.

import Control.Monad

summy :: Integer -> Integer -> [[Integer]]
summy k n = do
  ls <- replicateM k [1..n]
  guard (sum ls == n)
  return ls

Or just

summy k n = filter ((==n) . sum) $ replicateM k [1..n]

In the list monad, replicateM will generate all possible lists of length k composed of the numbers 1 .. n.

This does generate duplicates, like [1, 2, 1] and [1, 1, 2]. But so does your original methods.

Answer 2

As it happens, there's a lovely, efficient, and obscure (?) algorithm for enumerating the k-size partitions of n dating back to at least 1779. Donald Knuth -- who else? -- describes it in detail in a draft of the Art of Computer Programming, under Algorithm H. Here for your delectation is the algorithm in Haskell:

import Data.List (unfoldr)

partitions :: Int -> Int -> [[Int]]
partitions k n | k < 1 || k > n = []
partitions k n = initPartition : unfoldr (fmap (\x -> (x, x)) . nextPartition) initPartition
  where
    initPartition = (n-k+1) : replicate (k-1) 1

nextPartition :: [Int] -> Maybe [Int]
nextPartition [] = error "nextPartition should never be passed an empty list"
nextPartition [x] = Nothing
nextPartition (x:y:rest)
    | x-y > 1 = Just $ (x-1) : (y+1) : rest
    | otherwise = do
        (s, c, xs) <- go (x+y-1) rest
        Just $ (s-c) : c : xs
  where
    go _ [] = Nothing
    go s (z:zs)
        | x-z > 1 = let z' = z+1 in Just (s, z', z' : zs)
        | otherwise = do
            (s', c, zs') <- go (s+z) zs
            Just (s'-c, c, c:zs')