Dynamic programming from Algorithm Introduction (stations)

gistfile1.hs

-- Dynamic programming from Algorithm Introduction
module Stations where

import Control.Arrow
import Data.Array
import Data.Function.Memoize

data Line = Line {
    _size :: Int
  , _e1 :: Int
  , _e2 :: Int
  , _x1 :: Int
  , _x2 :: Int
  , _a1 :: Array Int Int
  , _a2 :: Array Int Int
  , _t1 :: Array Int Int
  , _t2 :: Array Int Int
}

lineExample :: Line
lineExample = Line {
    _size = 6
  , _e1 = 2
  , _e2 = 4
  , _x1 = 3
  , _x2 = 2
  , _a1 = listArray (1,6) [7,9,3,4,8,4]
  , _a2 = listArray (1,6) [8,5,6,4,5,7]
  , _t1 = listArray (1,5) [2,3,1,3,4]
  , _t2 = listArray (1,5) [2,1,2,2,1]
}

stations :: Line -> (Int, [Int])
stations ev = final
  where
    final
      | c1 <= c2  = r1
      | otherwise = r2
      where
        r1@(c1,_) = (***) (+x1) (1:) $ memof 1 size []
        r2@(c2,_) = (***) (+x2) (2:) $ memof 2 size []

    -- local memoized function does not leak tables
    -- FIXME: do we have to memoize the third argument?
    memof :: Int -> Int -> [Int] -> (Int,[Int])
    memof = memoize3 f

    f :: Int -> Int -> [Int] -> (Int,[Int])
    f 1 1 _ = (e1 + a1!1, [])
    f 2 1 _ = (e2 + a2!1, [])
    f 1 j ps
      | c1 <= c2  = r1
      | otherwise = r2
      where
        r1@(c1,_) = (***) (+ a1!j)              (1:) $ memof 1 (j-1) ps
        r2@(c2,_) = (***) (+ (t2!(j-1) + a1!j)) (2:) $ memof 2 (j-1) ps
    f 2 j ps
      | c1 <= c2  = r1
      | otherwise = r2
      where
        r1@(c1,_) = (***) (+ a2!j)              (2:) $ memof 2 (j-1) ps
        r2@(c2,_) = (***) (+ (t1!(j-1) + a2!j)) (1:) $ memof 1 (j-1) ps

    size = _size ev
    e1 = _e1 ev
    e2 = _e2 ev
    x1 = _x1 ev
    x2 = _x2 ev
    a1 = _a1 ev
    a2 = _a2 ev
    t1 = _t1 ev
    t2 = _t2 ev