Haskell による Ford-Fulkerson の実装

Posted on November 16, 2016

Ford-fulkerson 法による最大流問題のソルバを書いた. AOJ の Network Flow に投げる形で書いてある.

書いていて面倒だったところ

ソースコード

{-# LANGUAGE FlexibleContexts #-}
module Main where
import Data.List
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Control.Monad
import Data.Maybe
import Debug.Trace
import Data.Vector ((!))
import Control.Applicative

data Edge = Edge
   { to    :: Int -- 辺の貼られている先の頂点番号
   , revIx :: Int -- 辺に流せる最大流量へのインデックス
   , capIx :: Int -- 辺に流せる最大流量へのインデックス
   } deriving (Show,Eq)


type Graph = V.Vector (V.Vector Edge)

---------------
-- get input
--------------
getInts :: IO [Int]
getInts = unfoldr (B8.readInt . B8.dropWhile (==' ')) <$> BS.getLine

getGraph :: IO (Graph,MV.IOVector Int)
getGraph = do
  v:e:_ <- getInts
  cap <- MV.new (e*2)
  g <- MV.replicate v []
  forM_ [0..e-1] $ \i -> do
    u:v:c:_ <- getInts
    MV.write cap i c
    MV.write cap (i + e) 0
    us <- MV.read g u
    vs <- MV.read g v
    MV.write g u $ Edge v (e+i) i: us
    MV.write g v $ Edge u i (e+i): vs
    return ()
  g' <- V.map (V.fromList . reverse) <$> V.freeze g
  return (g', cap)

-----------------
-- graph printing
-----------------
printGraph :: Show a => V.Vector (V.Vector a) -> IO ()
printGraph = V.mapM_ (\i -> V.mapM (putStr . (' ':) . show) i >> putStrLn "")

printCap :: MV.IOVector Int -> IO ()
printCap v = print =<< V.freeze v


main :: IO ()
main = do
  (g,cap) <- getGraph
  -- printGraph g
  -- printCap cap
  r <- maxFlow 0 (V.length g - 1) g cap
  print r
  return ()


untilJustM :: Monad m => (a -> m (Maybe b)) -> V.Vector a -> m (Maybe b)
untilJustM act xs
  | V.null xs = return Nothing
  | otherwise = let v = V.head xs
                in do
      x <- act v
      case x of
        Nothing -> untilJustM act (V.tail xs)
        Just _ -> return x


maxFlow :: Int -> Int -> Graph -> MV.IOVector Int -> IO Int
maxFlow s t g cap = loop 0
  where
    loop f = do
      used <- MV.replicate (V.length g) False
      df <- dfs maxBound s t g cap used
      if df == 0
        then return f
        else loop (f + df)

dfs :: Int -> Int -> Int -> Graph
  -> MV.IOVector Int
  -> MV.IOVector Bool
  -> IO Int
dfs f v t g cap used
  -- | traceShow v False = undefined
  | v == t = return f
  | otherwise = do
      MV.write used v True
      fromMaybe 0 <$> untilJustM act (g ! v)
  where
    act :: Edge -> IO (Maybe Int)
    act e@Edge {to = to,capIx = ci, revIx = ri} = do
      c <- MV.read cap ci
      isUsed <- MV.read used to
      if c > 0 && not isUsed
      then do
        d <- dfs (min f c) to t g cap used
        if d > 0
        then do
          modifyCap cap e d
          return (Just d)
        else return Nothing
      else return Nothing

modifyCap :: MV.IOVector Int -> Edge -> Int -> IO ()
modifyCap cap Edge {capIx = ci, revIx = ri} d
  = do
  cv <- MV.read cap ci
  cr <- MV.read cap ri
  MV.write cap ci $ cv - d
  MV.write cap ri $ cr + d

comments powered by Disqus