Iterated updating of an IntMap

[emiruz]

Here is some code below. This is a dynamic programming version of a shortest path algorithm. The underlying grid this processes has maybe 10k cells and the paths function is particularly slow because it makes lots of modifications to an IntMap.

I’ve tried to improve performance by making only conditional updates, but its still epically slow.

Other than using a mutable vector – the ultimately solution – what other tactics are there which may improve the path function performance significantly?

#!/usr/bin/env runhaskell

import Data.List(elemIndex,foldl')
import Data.Char(ord)
import Text.Printf (printf)
import qualified Data.IntMap as M

main :: IO ()
main = do
  str <- getContents --readFile "sample.txt"
  let strip = filter(/='\n') str
      grid  = M.fromList $ zip [1..] $ strip
      ei c  = maybe 0 id $ elemIndex c $ strip
      (n,m) = (maybe 0 id $ elemIndex '\n' str, (length str) `div` (n+1))
      ps    = [(x,y) | x<-[1..m], y<-[1..n]]

      val x | x=='S' = -(ord 'a') | x=='E' = -(ord 'z') | otherwise = -(ord x)

      q g f (x,y)  = maybe 1000 f $ M.lookup ((x-1)*n+y) g
      q' g i = maybe 1000 id $ M.lookup i g

      adj (x,y) = filter (\p->v0-1 <= q grid val p) opts
        where v0 = q grid val (x,y)
              opts = map (\(a,b)->(x+a,b+y)) [(1,0),(-1,0),(0,1),(0,-1)]

      adj' = zip [1..] (map adj ps)
  
      start = M.insert (1+ei 'E') 0 $ M.fromList $ zip [1..] $ replicate (n*m) 1000

      paths vs | vs' == vs = vs | otherwise = paths vs'
        where vs' = foldl' (\acc (i,a)-> upd acc i $ ff (i,a)) vs adj'
              upd acc i (v0,v) = if v /= v0 then M.insert i v acc else acc
              ff (i,a) = let v0 = q' vs i in (v0, minimum $ v0:(map ((+1).(q vs id)) a))

      solve = paths start
      part1 = maybe 0 id $ M.lookup (1+ei 'S') solve
      part2 = minimum [x | (x,c) <- zip (M.elems solve) $ strip, elem c "Sa"]

  printf "Part 1: %d, Part 2: %d\n" part1 (part2 :: Int)

[jaror]

• The first thing I would try is compiling it with optimizations (ghc -O2) instead of running it interpreted with runhaskell.

• I notice that you do an expensive equality check on intmaps in paths, instead you could simply directly track if something has changed. Something like this:

  data StrictTuple a b = T2 !a !b
  

      paths vs
        | hasChanged = paths vs'
        | otherwise = vs
        where T2 vs' hasChanged = foldl' (\acc (i,a) -> upd acc i $ ff (i,a)) (T2 vs False) adj'
              upd (T2 acc hasChanged) i (v0,v) = if v /= v0 then T2 (M.insert i v acc) _ else T2 acc hasChanged
              ff (i,a) = let v0 = q' vs i in (v0, minimum $ v0:(map ((+1).(q vs id)) a))
  

  I don’t quite know how to best fill in that hole. You’d need to know whether the insert has actually changed anything. You could consider using:

  updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a, IntMap a)
  

  Or maybe it is guaranteed that anything you insert will be different from what was already there? Then you can just use True in the place of that hole.

• Also, you can use a nested immutable vector for adj. Like change the type from [(Int, [(Int, Int)])] to Vector (Unboxed.Vector (Int, Int))

[emiruz]

@jaror

Also, you can use a nested immutable vector for adj. Like change the type from [(Int, [(Int, Int)])] to Vector (Unboxed.Vector (Int, Int))

Yes a good point. This would result in O(1) lookups I presume? There are at least 41x171=7011 lookups (grid size) on each iteration; not sure how many iterations. I think IntMap is O(log N) for lookups, so it would save time by a factor of roughly log(7011) ~ 9.

@jaror

I notice that you do an expensive equality check on intmaps in paths, instead you could simply directly track if something has changed. Something like this:

I check if old/new values are the same as part of the updates, so as indicated I should be able to bubble that up somehow, although if there aren’t many iterations I’m not sure it’ll save much time.

[emiruz]

@emiruz

ff (i,a) = let v0 = q' vs i in (v0, minimum $ v0:(map ((+1).(q vs id)) a))

I see now that I can avoid q' vs i by directly passing into the foldl . I.e. instead just adj' I can pass zip vs adj and then i won’t need to do the look up. That saves another 7011 x 9 cycles per iteration.

Huh … and adj’ is in the same order anyway, so there is no need for me to do lookups in the first place …

Yeah, I think I’m good to go. Thank you @jaror for enacting the change in perspective in me.

[emiruz]

Here is the updated snippet for the paths function:

paths vs | chg==0 = vs | otherwise = paths vs'
   where (chg, vs') = foldl' (\acc y@(v0,(i,_))-> upd acc i (ff y) v0) (0,vs) (zip (M.elems vs) adj')
         upd (chg, acc) i v v0 = if v /= v0 then (1, M.insert i v acc) else (chg, acc)
         ff (v0,(_,a)) = minimum $ v0:(map ((+1).(q vs id)) a)

ghc -O2 made a huge difference. drop q' vs i by zipping it into the fold also made a big difference. Removing vs==vs’ seemed to make little difference. All together, it took the compiled run time from 2.2 sec to 0.7 sec on a 2.8Ghz i7.

[jaror]

I think you should be careful with using (lazy) tuples as the accumulator of the foldl'. Maybe -O2 is able to catch it but you should probably at least use bang patterns:

paths vs | chg==0 = vs | otherwise = paths vs'
   where (!chg, !vs') = foldl' (\acc y@(v0,(i,_))-> upd acc i (ff y) v0) (0,vs) (zip (M.elems vs) adj')
         upd (!chg, !acc) i v v0 = if v /= v0 then (1, M.insert i v acc) else (chg, acc)
         ff (v0,(_,a)) = minimum $ v0:(map ((+1).(q vs id)) a)

[jaror]

And you could also check if switching to strict IntMap helps:

import qualified Data.IntMap.Strict as M

[tomjaguarpaw]

I strongly recommend the Strict version of IntMap. Regarding the tuple accumulator, I recommend strict-wrapper: Lightweight strict types.

For explanations, see my articles

• nested-strict-data
• make-invalid-laziness-unrepresentable

[emiruz]

Doesn’t seem to make a difference. foldl and foldl’ make a small difference.