Haskell で Simpath を実装してみる

ふと[以前に書いた記事](http://qiita.com/yasuabe1362/items/78d61e09cb791d9f9154 「オフラインリアルタイムどう書く第四回の参考問題」を解いてみた)を思い出した。５年近くも前のグリッドの経路の数を数えるコードで、シンプルだけど少しグリッドのサイズが増えるとすぐに組み合わせ爆発してしまう。

Haskell を書くのも数年ぶりだけど、今では simpath もだいたいわかるので、書き直してみようと思う。

辺と頂点

グラフの辺を Edge 型、頂点を Node 型として表す.
Simpath/Edge.hs

〜冒頭略〜

type Node = Int
data Edge = Edge { left :: Node, right :: Node } deriving (Show, Eq, Ord)

edge :: Node -> Node -> Edge
edge  = (liftA2 . liftA2) Edge min max

start :: Edge
start = edge 0 1

onBoth :: (a -> a -> b) -> (Node -> a) -> Edge -> b
onBoth f g (Edge l r) = on f g l r

either:: (Node -> Bool) -> Edge -> Bool
either = onBoth (||)

contains :: Node -> Edge -> Bool
contains = either . (==)

modify :: (Node -> Node) -> Edge -> Edge
modify = onBoth edge

opposite :: Node -> Edge -> Node
opposite n (Edge l r) = if l == n then r else l

isOpen :: Edge -> Bool
isOpen = onBoth (/=) id

find :: Node -> Set Edge -> Maybe Edge
find = Foldable.find . contains

connect :: Node -> Node -> Edge -> Edge
connect from to = edge from . opposite to

フロンティア

これが Simpath の要点となるデータ。一度訪れていて後でまた訪れうる頂点の集合を表す。この実装では、訪れうるが訪れるとパスが無効になる頂点も合わせて保持している。

Simpath/Frontier.hs

〜冒頭略〜

type Used  = Set Node
type Edges = Set Edge
data Frontier = Frontier { edges :: Edges, used :: Used } deriving (Show, Eq, Ord)

initial :: Frontier
initial = Frontier (Set.singleton E.start) (Set.empty)

modify :: (Edges -> Edges) -> (Used -> Used) -> Frontier -> Frontier
modify f g = Frontier <$> f . edges <*> g . used

add :: Edge -> Frontier -> Maybe Frontier
add e@(E.Edge l r) fr@(Frontier _ used) =
    justIf (notUsed e && E.isOpen e') $ modify (Set.insert e') id fr'
  where (e', fr')     = runState (connect l r >>= connect r . E.opposite r) fr
        notUsed       = not . E.either isUsed where isUsed = flip Set.member used
        connect n1 n2 = state (mapOrElse <$> step <*> (,)(E.edge n1 n2) <*> find)
          where find      = E.find n1 . edges
                step fr e = (E.connect n2 n1 e, update fr)
                  where update = modify (Set.delete e) (Set.insert n1)

proceed :: Border -> Frontier -> (Maybe Frontier, Maybe Frontier)
proceed (B.Border edge done) = (,) <$> proceedHi <*> proceedLo
  where proceedHi = add edge >=> proceedLo
        proceedLo = mapOrElse (\d -> justIf <$> not . contains d <*> removeUsed d) Just $ done
          where removeUsed = modify id . Set.delete
                contains n = Foldable.any (E.contains n) . edges

Border

辺と、その辺を処理すると「その後は再到達し得ない頂点」のペアを保持するデータ型。Edge のリストを順に処理することで Border のリストを算出することができる。

Simpath/Border.hs

〜冒頭略〜

data Border = Border { edge :: Edge, done :: Maybe Node } deriving (Show)

borders :: [Edge] -> [Border]
borders es = reverse $ evalState bordersState Set.empty
  where
    bordersState = sequence $ map fromEdge $ reverse es
    fromEdge :: Edge -> State (Set Node) Border
    fromEdge e@(Edge l _) = do used <- get
                               put $ Set.insert l used
                               return $ Border e $ justIf (Set.notMember l used) l

CounterMap

ある Frontier に到達するパスが何通りあるか保持するデータ型。countPaths 関数で全ての Border を超えて最終的に合算された経路数を返す。（ネットで見つけた他の実装では、一旦縮約したグラフを作ってから、辿り直して合算する方式のものもあったが、ここでは Frontier を最終地点まで進めると同時に、経路数もついでに足し合わせていくようにした。）

Simpath/CounterMap.hs

〜冒頭略〜

type CounterMap = Map Frontier Integer

merge :: Integer -> CounterMap -> Frontier -> CounterMap
merge cnt cm fr = Map.insertWith (+) fr cnt cm

proceedAll :: Border -> CounterMap -> CounterMap
proceedAll b = Foldable.foldl f Map.empty . Map.toList
  where f cm (fr, cnt) = on (.) recount lo hi cm
          where (hi, lo) = F.proceed b fr
                recount  = flip $ Foldable.foldl $ merge cnt

countPaths :: [Border] -> Integer
countPaths = headCount . foldl (flip proceedAll) initialMap
  where initialMap = Map.singleton F.initial 1
        headCount  = head . Map.elems

共通コード

Simpath/Common.hs

module Simpath.Common where

justIf :: Bool -> a -> Maybe a
justIf b a = if b then Just a else Nothing

mapOrElse :: (a -> b) -> b -> Maybe a -> b
mapOrElse f b ma = case ma of { Just a -> f a; _ -> b }

実行コード

実行引数で与えられたサイズのグリッドを作って、パスを数えるコード。グリッドの辺の集合は前半部分を作れば、それを裏返す感じで残り半分も作ることができる。

Main.hs

〜冒頭略〜

gridEdges :: Int -> [Edge]
gridEdges size = upper ++ lower
  where
    upper = snd $ foldl (\(c, ts) n -> (c + n, ts ++ edgesAt n c)) (0, []) [1 .. size-1]
      where edgesAt n acc = map (+ acc) [1 .. n] >>= addPair
              where addPair i = let f = edge i . (i + n +) in map f [0, 1]
    lower = reverse $ map (modify (size^2 + 1 -)) upper

main :: IO ()
main = do args <- getArgs
          print $ countPaths $ borders $ gridEdges $ (read $ head args :: Int)

結果

前にやったとき２５秒かかっていたサイズ 6 が以下の通り。

$ ghc -O Main.hs; time ./Main 6
1262816
./Main 6  0.02s user 0.01s system 49% cpu 0.070 total

お姉さんのスーパーコンピューターで２５万年かかっていたサイズ11 (10×10) でも約2分で終わる。

$ ghc -O Main.hs; time ./Main 11
1568758030464750013214100
./Main 11  120.25s user 34.91s system 87% cpu 2:58.15 total

参考

『フカシギの数え方』 おねえさんといっしょ！ みんなで数えてみよう！