• + 0 comments

    DP with Memoization in Haskell

    import Data.Array (Ix, array, bounds, inRange, (!))
    import Prelude hiding (Right)
    
    data Direction = Neutral | Down | Right deriving (Eq, Ord, Ix)
    
    paths :: Int -> Int -> Int -> Integer
    paths m n k = go (1, 1, k, Neutral)
      where go (i, j, k, dir)
              | i == m && j == n = 1
              | otherwise = case dir of
                Neutral -> mod (get (i+1, j, k, Down) + get (i, j+1, k, Right)) (10^9 + 7)
                Down -> mod (get (i+1, j, k, Down) + get (i, j+1, k-1, Right)) (10^9 + 7)
                Right -> mod (get (i+1, j, k-1, Down) + get (i, j+1, k, Right)) (10^9 + 7)
            a = array ((1, 1, 0, Down), (m, n, k, Right))
                [(c, go c) | c <- (,,,) <$> [1..m] <*> [1..n] <*> [0..k] <*> [Down, Right]]
            get x | inRange (bounds a) x = a ! x
                  | otherwise = 0
                  
    helper :: [String] -> Integer 
    helper arr = do
        let new = map (read :: String -> Int) arr
        paths (new!!0) (new!!1) (new!!2)
        
    main :: IO()
    main = do
        raw <- getContents
        mapM_ print $ map helper $ tail $ map words $ lines raw