Search code examples
haskellprefix-sum

Undoing a prefix sum


The simple way to calculate a prefix sum in Haskell is

scanl1 (+) [1,2,3,4,5,6]

which produces the output

[1,3,6,10,15,21]

I needed to write the inverse, and this is what I came up with:

undo_prefix_sum :: (Num a) => [a] -> [a]
undo_prefix_sum s = init $ snd $ 
    foldr (\cur (tot, l) -> (cur, (tot-cur):l)) 
          (last s, []) 
          (0:s)

This seems to be correct (but I might have missed something). Is there a simpler or more efficient way to do this, possibly using a scan?


Solution

  • In my opinion, as prefix_sum is naturally expressed as a fold, undo_prefix_sum is more naturally expressed as an unfold.

    import Data.List
    
    undo_prefix_sum  = unfoldr (\xs -> if null xs 
                                       then Nothing
                                       else let (h:t) = xs 
                                            in Just (h,  map (subtract h) t))
    

    unfoldr uses a function to build a list starting from a seed value. In this case, the seed is a list itself; we use the first element of the seed in our result, and (a modified version of) the rest of the list as the seed to compute the rest of the result recursively.