Submission #34031960

---

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.Unboxed.Mutable as MUV

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
-- 0番が発電所、1～N 個は都市
    uf <- newUF (succ n)
-- 切れない電線を配線する
    sequence_ [uniteUF uf u v | ((u,v),True) <- zip (elems uvA) (elems uvX)]
-- 逆順に Xi を繋ぎつつ、都市数を追跡
    ansS <- mapM (action uf) (map (uvA !) $ reverse xs)
    return $ reverse ansS
  where
-- 電線の配列（発電所は0に統合）
    uvA = listArray (1,e) [(pp u, pp v) | (u:v:_) <- uvs]
    pp x
      | x <= n    = x
      | otherwise = 0
-- 電線が切れないかチェック表
    uvX = accumArray (&&) True (1,e) [(x,False) | x <- xs]

action uf (u, v) =
  do
    pp <- getRoot uf 0
    a <- MUV.read uf pp
    uniteUF uf u v
    return (pred $ negate a)

-- UnionFind
type UnionFind = MUV.IOVector Int

newUF :: Int -> IO UnionFind
newUF n = MUV.replicate n (-1)

getRoot :: UnionFind -> Int -> IO Int
getRoot vec i = loop i []
  where
    loop :: Int -> [Int] -> IO Int
    loop i ks =
      do
        k <- MUV.read vec i
        if k < 0 then
          do
            forM_ ks (\k -> do MUV.write vec k i)   -- 経路圧縮
            return i
          else
            loop k (i:ks)

uniteUF :: UnionFind -> Int -> Int -> IO Bool
uniteUF vec i j =
  do
    a <- getRoot vec i
    b <- getRoot vec j
    if a == b then return False else
      do
        r <- MUV.read vec a
        s <- MUV.read vec b
        if r < s then
          do
            MUV.write vec a b
            MUV.write vec b (r+s)
          else do
            MUV.write vec b a
            MUV.write vec a (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