Data.Function.fix で loop 制御

競プロ典型 90 問 40 日目 - Get More Money（★7)

import Control.Monad
import Control.Monad.ST
import Control.Monad.Primitive
import qualified Data.ByteString.Char8 as C
import Data.List
import Data.Function
import Data.Functor
import Data.Sequence (Seq (Empty,(:<|)),(|>))
import qualified Data.Sequence as S
import Data.Vector.Generic ((!),(//),empty,snoc)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

main = ints >>= \(n:w:_) -> sol n w <$> ints <*> replicateM n ints >>= print

ints = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine

sol n w as css = sum as-runST (dinic g 0 (n+1))
  where
  es = concatMap (\(i,a) -> [(0,i,w),(i,n+1,a)]) $ zip [1..] as
  es' = map (\(u,v) -> (u,v,big)) . concat $ zipWith ((. tail) . map . (,)) [1..] css
  g = foldl' addEdge (G.replicate (n+2) empty) (es++es') :: Graph

big = 10^10 :: Int
type Graph = V.Vector (U.Vector (Int, Int, Int))
type MGraph m = V.Vector (UM.MVector (PrimState m) (Int, Int, Int))

addEdge g (v,u,c) = do
  let
    vs = g!v
    us = g!u
    !vs' = snoc vs (u,c,G.length us)
    !us' = snoc us (v,0,G.length vs)
  g//[(v,vs'),(u,us')]

dinic g s t = do
  g' <- G.mapM G.thaw g
  lvl <- UM.new (G.length g)
  UM.set lvl 0
  flip fix 0 $ \outloop fl -> do
    dv <- bfs g' s
    if dv!t == -1
      then return fl
      else do
        UM.set lvl 0
        flip fix fl $ \inloop f -> do
          f' <- dfs g' dv lvl s t big
          if f'>0
            then inloop (f+f')
            else outloop f

bfs g s = do
  dv <- UM.new (length g)
  UM.set dv (-1 :: Int)
  UM.write dv s 0
  aux (S.singleton s) dv
  U.freeze dv
  where
  aux Empty _     = return ()
  aux (v:<|q) dv  = do
    d <- UM.read dv v
    q' <- U.freeze (g!v)
      >>= U.filterM (\(u,c,_) -> UM.read dv u <&> (c>0 &&) . (-1==))
      >>= U.foldM' (\q (u,_,_) -> UM.write dv u (d+1) >> return (q|>u)) q
    aux q' dv

dfs g dv lvl v t fl
  | v==t = return fl
  | otherwise =
    fix $ \loop -> do
      let edges = g!v
      i <- UM.read lvl v
      if i<UM.length edges
        then do
          UM.write lvl v (i+1)
          (to,c,r) <- UM.read edges i
          if c>0 && dv!v<dv!to
            then do
              f <- dfs g dv lvl to t (min fl c)
              if f>0
                then do
                  UM.write edges i (to,c-f,r)
                  UM.modify (g!to) (\(to',c',r') -> (to',c'+f,r')) r
                  return f
                else loop
            else loop
        else return 0