パスを数える simpath algorithm を実装した (Haskell)

パスを数える simpath algorithm を実装した。knuth先生のあの本にのっているとか。

参考動画 http://youtu.be/Q4gTV4r0zRs

simpathは動的計画法(DP)の一種。ただし、そのまま DPの表を作るとメモリ不足とかで頓死するので、動的に表を作成し、不要部分は削除、共有可能な部分は共有するという方針。

それでも、メモリ消費は大きくてHaskellで実装した下のコードでは 9×9 は計算できたけど 10×10 はメモリが足りなくて実行できなかった。工夫すればできそうだけど、そもそもHaskellで書くのがまちが(ry

{-# LANGUAGE TupleSections #-}

import Data.IntMap  (IntMap, insert, insertWith, delete, empty, fromList)
import Data.Map (fromListWith, toList)
import qualified Data.IntMap as IntMap
import Data.Maybe (fromJust, mapMaybe)
import System.Environment (getArgs)

type Node = Int
type Edge = (Node, Node)
type ExEdge = (Edge, (Bool, Bool))
type Mate = IntMap Node
type State = (Mate, Integer)

addEdge :: Node -> Node -> Mate -> Maybe Mate
addEdge u v m
    | a == 0 || b == 0 = Nothing -- split
    | a == u && b == v = Just m1 -- u-v
    | a == u && b /= v = Just m2 -- u-v-b
    | a /= u && b == v = Just m3 -- a-u-v
    | a == v && b == u = Nothing -- cycle
    | a == b           = Nothing -- cycle
    | otherwise        = Just m4 -- a-u-v-b
    where [a,b] = map (`lookup'` m) [u,v] -- a-u v-b
          [m1,m2,m3,m4] = map (foldr (uncurry insert) m)
                          [[(u,v), (v,u)],
                           [(u,b), (v,0), (b,u)],
                           [(v,a), (u,0), (a,v)],
                           [(u,0), (v,0), (a,b), (b,a)]]

lookup' :: Node -> IntMap a -> a
lookup' k m = fromJust $ IntMap.lookup k m

nextMate :: ExEdge -> Mate -> [Mate]
nextMate ((u,v),(eu,ev)) m = mapMaybe ((exCheck u eu =<<).(exCheck v ev =<<)) [addEdge u v m', Just m']
    where m' = weakInsert u u.weakInsert v v $ m
          weakInsert = insertWith (flip const)

exCheck :: Node -> Bool -> Mate -> Maybe Mate
exCheck _ False m = Just m
exCheck x True m | y == x || y == 0 = Just $ delete x m
                 | otherwise        = Nothing -- break
                 where y = lookup' x m

step :: [State] -> ExEdge -> [State]
step ss e = toList.fromListWith (+) $ concatMap nextState ss
  where nextState (m,c) = map (, c) $ nextMate e m

toExList :: [Edge] -> [ExEdge]
toExList [] = []
toExList ((u,v):es) = ((u,v),(notElem u es', notElem v es')) : toExList es
    where es' = concatMap (\(a,b) -> [a,b]) es

count :: [Edge] -> Maybe Integer
count es = lookup s.foldl step [(empty, 1)] $ exs
  where exs = init.toExList $ es ++ [(1,m)]
        m = maximum $ map (uncurry max) es
        s = fromList [(1,m),(m,1)]

gridEdges :: Int -> [Edge]
gridEdges n = concat [rows i ++ cols i| i <- [1..n]] ++ rows (n+1)
    where pos i j = (i-1) * (n+1) + j
          row i j = (pos i j, pos i (j+1))
          col i j = (pos i j, pos (i+1) j)
          rows i = [row i j | j <- [1..n]]
          cols i = [col i j | j <- [1..n+1]]

main :: IO ()
main = print.count.gridEdges.read.head =<< getArgs