cojna/iota ライブラリは最強

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

{-# LANGUAGE LambdaCase, RecordWildCards #-}

import Control.Arrow
import Control.Monad
import           Control.Monad.ST
import qualified Data.ByteString.Char8             as C
import Data.Function
import Data.List
import qualified Data.Map.Strict                   as M
import Data.Maybe
import Data.Vector.Unboxed ((!))
import qualified Data.Vector.Fusion.Stream.Monadic as MS
import qualified Data.Vector.Unboxed               as U
import qualified Data.Vector.Unboxed.Mutable       as UM

-- cojna / iota library
import Data.Graph.BipartiteMatching

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 n t as = r >>= Just . zipWith dir [0..] . U.toList
  where
  (xs,ys) = U.splitAt n . U.fromList $ map (\[x,y] -> (x,y)) as
  mp = U.foldl' (\m (v,k) -> M.insert k v m) M.empty $ U.indexed ys
  dt = U.fromListN 8 $ map ((* t) *** (* t)) [(1,0),(1,1),(0,1),(-1,1),(-1,0),(-1,-1),(0,-1),(1,-1)]
  mkEdge i k = (i,) . (+ n) <$> M.lookup (x+dx,y+dy) mp
    where 
    (x,y) = xs!i
    (dx,dy) = dt!k
  es = U.fromList . catMaybes $ mkEdge <$> [0..n-1] <*> [0..7]
  r = matching n es
  dir i j = fromJust $ find ((==(x'-x,y'-y)) . (dt!) . pred) [1..8]
    where
    (x,y) = xs!i
    (x',y') = ys!(j-n)

matching n es = runST $ do
  builder <- newBipartiteMatchingBuilder (2*n)
  U.forM_ es $ uncurry (addEdgeBMB builder)
  bm <- buildBipartiteMatching builder
  matchCnt <- runBipartiteMatching bm
  if matchCnt<n
    then pure Nothing
    else do
      ds <- UM.new n
      forM_ [0..n-1] $ \i -> UM.read (matchBM bm) i >>= UM.write ds i
      Just <$> U.freeze ds