Search code examples
haskellrecursionstate-monadaccumulator

Recursive state monad for accumulating a value while building a list?


I'm totally new to Haskell so apologies if the question is silly.

What I want to do is recursively build a list while at the same time building up an accumulated value based on the recursive calls. This is for a problem I'm doing for a Coursera course, so I won't post the exact problem but something analogous.

Say for example I wanted to take a list of ints and double each one (ignoring for the purpose of the example that I could just use map), but I also wanted to count up how many times the number '5' appears in the list.

So to do the doubling I could do this:

foo []     = []
foo (x:xs) = x * 2 : foo xs

So far so easy. But how can I also maintain a count of how many times x is a five? The best solution I've got is to use an explicit accumulator like this, which I don't like as it reverses the list, so you need to do a reverse at the end:

foo total acc []     = (total, reverse acc)
foo total acc (x:xs) = foo (if x == 5 then total + 1 else total) (x*2 : acc) xs

But I feel like this should be able to be handled nicer by the State monad, which I haven't used before, but when I try to construct a function that will fit the pattern I've seen I get stuck because of the recursive call to foo. Is there a nicer way to do this?

EDIT: I need this to work for very long lists, so any recursive calls need to be tail-recursive too. (The example I have here manages to be tail-recursive thanks to Haskell's 'tail recursion modulo cons').


Solution

  • This is a simple fold

    foo :: [Integer] -> ([Integer], Int)
    foo []       = ([], 0)
    foo (x : xs) = let (rs, n) = foo xs
                    in (2 * x : rs, if x == 5 then n + 1 else n)
    

    or expressed using foldr

    foo' :: [Integer] -> ([Integer], Int)
    foo' = foldr f ([], 0)
      where
        f x (rs, n) = (2 * x : rs, if x == 5 then n + 1 else n)
    

    The accumulated value is a pair of both the operations.

    Notes:

    • Have a look at Beautiful folding. It shows a nice way how to make such computations composable.
    • You can use State for the same thing as well, by viewing each element as a stateful computation. This is a bit overkill, but certainly possible. In fact, any fold can be expressed as a sequence of State computations:

      import Control.Monad
      import Control.Monad.State
      
      -- I used a slightly non-standard signature for a left fold
      -- for simplicity.
      foldl' :: (b -> a -> a) -> a -> [b] -> a
      foldl' f z xs = execState (mapM_ (modify . f) xs) z
      

      Function mapM_ first maps each element of xs to a stateful computation by modify . f :: b -> State a (). Then it combines a list of such computations into one of type State a () (it discards the results of the monadic computations, just keeps the effects). Finally we run this stateful computation on z.