StateT and monad non-determinism: a simple example
As part of learning how to work with StateT and the non-determinism monad, I would like to write a function that uses them to enumerate sections of an integer (while using integers at the same time). For example, passing an argument 4
should result in [[1,1,1,1],[1,1,2],[2,2],[1,3],[4]]
(uniqueness doesn't matter, I'm only more interested in getting to working code).
(Also, I know there is a recursive solution for creating partitions as well as dynamic programming and creating function-based solutions for partition counting - the goal of this exercise is to create a minimal working example combining StateT and [].)
Here's my attempt, which was designed to work with any input less than or equal to 5:
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wall #-}
import CorePrelude
import Control.Monad.State.Lazy
sumState :: StateT Int [] [Int]
sumState = do
m <- lift [1..5]
n <- get <* modify (-m+)
case compare n 0 of
LT -> mzero
EQ -> return [m]
GT -> fmap (n:) sumState
runner :: Int -> [([Int],Int)]
runner = runStateT sumState
I am using runStateT
and not evalStateT
to help with debugging (useful to see final state values). As I said, I am not too worried about creating unique partitions, as at first I would just like to understand the correct way to use these two monads.
Loading in GHCi
and evaluating runner 4
results in the following, and I am confused as to why the above code produces this output.
[([4,3,2,1,1],-1),([4,3,2,1,2],-2),([4,3,2,1,3],-3),([4,3,2,1,4],-4),([4,3,2,1,5],-5),([4,3,2,1],-1),([4,3,2,2],-2),([4,3,2,3],-3),([4,3,2,4],-4),([4,3,2,5],-5),([4,3,1,1],-1),([4,3,1,2],-2),([4,3,1,3],-3),([4,3,1,4],-4),([4,3,1,5],-5),([4,3,1],-1),([4,3,2],-2),([4,3,3],-3),([4,3,4],-4),([4,3,5],-5),([4,2,1,1],-1),([4,2,1,2],-2),([4,2,1,3],-3),([4,2,1,4],-4),([4,2,1,5],-5),([4,2,1],-1),([4,2,2],-2),([4,2,3],-3),([4,2,4],-4),([4,2,5],-5),([4,1,1],-1),([4,1,2],-2),([4,1,3],-3),([4,1,4],-4),([4,1,5],-5),([4,1],-1),([4,2],-2),([4,3],-3),([4,4],-4),([4,5],-5)]
What am I doing wrong? What's the correct way to combine StateT and [] to enumerate sections?
source to share
You only have two errors. The first one is here:
n <- get <* modify (-m+)
This gets the value n
before it is subtracted m
. You almost certainly want
n <- modify (-m+) >> get
or
modify (-m+)
n <- get
if you prefer this spelling. Another is that you put the current state in the list instead of the value you add to the branch GT
:
GT -> fmap (n:) sumState
Change this to
GT -> fmap (m:) sumState
and you're golden:
*Main> runner 4
[([1,1,1,1],0),([1,1,2],0),([1,2,1],0),([1,3],0),([2,1,1],0),([2,2],0),([3,1],0),([4],0)]
source to share