Vector of Vector は General.mapM で解凍/凍結

競プロ典型 90 問 77 日目 - Planes on a 2D Plane（★7）

import Control.Arrow
import Control.Monad
import Control.Monad.ST
import Control.Monad.Primitive
import qualified Data.ByteString.Char8 as C
import Data.Function
import Data.Functor
import Data.List
import Data.Maybe
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,t] -> replicateM (2*n) ints >>= fmt . sol n t

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

fmt Nothing   = putStrLn "No"
fmt (Just ds) = putStrLn "Yes" >> (putStrLn . unwords $ map show ds)

sol :: Int -> Int -> [[Int]] -> Maybe [Int]
sol n t as = if fl<n then Nothing else ds
  where
  es = concatMap (\i -> [(0,i,1),(i+n,2*n+1,1)]) [1..n]
  xy = U.fromListN (2*n+2) $ (-1,-1):map (\[x,y] -> (x,y)) as ++ [(10^10+1,10^10+1)]
  dt = U.fromListN 8 $ map ((* t) *** (* t)) [(1,0),(1,1),(0,1),(-1,1),(-1,0),(-1,-1),(0,-1),(1,-1)]
  es' = catMaybes $ mkEdge <$> [1..n] <*> [n+1..2*n]
  mkEdge i j = if U.any (==(x'-x,y'-y)) dt
    then Just (i,j,1)
    else Nothing
    where
    (x,y) = xy!i
    (x',y') = xy!j
  g = foldl' addEdge (V.replicate (2*n+2) empty) (es++es') :: Graph
  (fl,g') = runST $ dinic g 0 (2*n+1)
  ds = Just $ mapMaybe (dir g') [1..n]
  dir g i = find ((==(x'-x,y'-y)) . (dt!) . pred) [1..8]
    where
    (x,y) = xy!i
    (x',y') = xy!j
    j = maybe 0 (\(j,_,_) -> j) $ U.find (\(j,c,_) -> n<j && j<=2*n && c==0) (g!i)

type Graph = V.Vector (U.Vector (Int, Int, Int))

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

-- dinic :: PrimMonad m => Graph -> Int -> Int -> m (Int, Graph)
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 fmap (fl,) (G.mapM G.freeze g')
      else do
        UM.set lvl 0
        flip fix fl $ \inloop f -> do 
          f' <- dfs g' dv lvl s t (10^10+1)
          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