Submission #34030932

---

Source Code Expand

import Control.Monad
import qualified Data.ByteString.Char8 as BS
import Data.Char
import Data.List

import Data.Array
import Control.Applicative
import qualified Data.Vector.Mutable as MV

main = do
  [n,m,e] <- bsGetLnInts
  uvs <- replicateM e bsGetLnInts
  [q] <- bsGetLnInts
  xs <- replicateM q (head <$> bsGetLnInts)
  ans <- abc264e n m e uvs q xs
  mapM_ print ans

bsGetLnInts :: IO [Int]
bsGetLnInts = BS.getLine >>= return . unfoldr (BS.readInt . BS.dropWhile isSpace)

abc264e :: Int -> Int -> Int -> [[Int]] -> Int -> [Int] -> IO [Int]
abc264e n m e uvs q xs =
  do
-- 先頭 N 個は都市、後半 M 個が発電所
    uf <- newUF $ replicate n (1,False) ++ replicate m (0, True)
-- 切れない電線を配線する
    sequence_ [uniteUF uf u v | ((u,v),True) <- zip (elems uvA) (elems uvX)]
-- 最後まで通電している都市の数
    ansQ <- getPowered uf
-- 逆順に Xi を繋ぎつつ、都市数を追跡
    foldM (action uf) [ansQ] (map (uvA !) $ reverse $ tail xs)
  where
-- 電線の配列
    uvA = listArray (1,e) [(pred u, pred v) | (u:v:_) <- uvs]
-- 電線が切れないかチェック表
    uvX = accumArray (&&) True (1,e) [(x,False) | x <- xs]

action uf (acc:accs) (u, v) =
  do
    (a, (cc1,pp1)) <- getRoot uf u
    (b, (cc2,pp2)) <- getRoot uf v
    if a == b then return (acc:acc:accs) else do
      uniteUF uf a b
      return $ case (pp1, pp2) of
        (True, False) -> acc+cc2:acc:accs
        (False, True) -> acc+cc1:acc:accs
        _             -> acc    :acc:accs

{- foldl : Since: 0.12.3.0
getPowered uf = MV.foldl step 0 uf
  where
    step acc (Right (cc,True)) = cc + acc
    step acc _ = acc
-}
getPowered uf = do
  foldM (\acc i -> do
    x <- MV.read uf i
    return $ case x of
      Right (cc,True) -> acc + cc
      _ -> acc
    ) 0 [0 .. pred $ MV.length uf]

uniteUF = mkUniteUF (\(cc1,pp1) (cc2,pp2) -> (cc1 + cc2, pp1 || pp2))

-- UnionFind

-- @gotoki_no_joe
type UnionFind a = MV.IOVector (Either Int a)

newUF :: [a] -> IO (UnionFind a)
newUF xs = do
  v <- MV.new (length xs)
  forM_ (zip [0..] xs) (\(i,x) -> MV.write v i (Right x))
  return v

getRoot :: UnionFind a -> Int -> IO (Int, a)
getRoot vec i = loop vec i []
  where
    loop vec i ks =
      do
        k <- MV.read vec i
        case k of
          Right x -> do
            forM_ ks (\k -> MV.write vec k (Left i))
            return (i, x)
          Left j -> loop vec j (i:ks)

mkUniteUF :: (a -> a -> a) -> UnionFind a -> Int -> Int -> IO Bool
mkUniteUF f vec i j =
  do
    (a, r) <- getRoot vec i
    (b, s) <- getRoot vec j
    if a == b then return False else
      do
        MV.write vec a (Left b)
        MV.write vec b (Right $ f r s)
        return True

Submission Info

Compile Error

Loaded package environment from /home/contestant/.ghc/x86_64-linux-8.8.3/environments/default

Judge Result