Submission #34034150

---

Source Code Expand

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

import Data.Array
import qualified Data.IntMap as IM
import Control.Applicative

main = do
  [n,m,e] <- bsGetLnInts
  uvs <- replicateM e bsGetLnInts
  [q] <- bsGetLnInts
  xs <- replicateM q (head <$> bsGetLnInts)
  let 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] -> [Int]
abc264e n m e uvs q xs = ccs
  where
-- 0番が発電所、1～N 個は都市
    uf0 = newUF (succ n)
-- 電線の配列
    uvA = listArray (1,e) [(pp u, pp v) | (u:v:_) <- uvs]
    pp x = if n < x then 0 else x
-- 電線が切れないかチェック表
    uvX = accumArray (&&) True (1,e) [(x,False) | x <- xs]
-- 切れない電線を配線する
    uf1 = foldl' {-'-} uniteUF uf0 [uv | (uv,True) <- zip (elems uvA) (elems uvX)]
-- 逆順に Xi を繋いだufを作りつつ、都市数を追跡
    (_,ccs) = mapAccumR step uf1 xs
    step uf xi = (uniteUF uf (uvA ! xi), pred cc1)
      where
        (pp, cc1) = getRoot uf 0

-- UnionFind
-- 編集が定数時間でできるIntMapを使ってみる

type UnionFind = IM.IntMap Int
-- 要素は 親id または 分割の要素数(負)

-- 0からNまでのN要素が独立した初期状態を作る
newUF :: Int -> UnionFind
newUF n = IM.fromAscList [(i,-1) | i <- [0..n]]

-- 補助関数 ノードの根まで辿り、代表idと分割の要素数を返す
getRoot :: UnionFind -> Int -> (Int,Int)
getRoot uf i = loop i
  where
    loop i
      | k >= 0 = loop k
      | True   = (i, negate k)
      where
        k = uf IM.! i

-- ふたつのノードが同じ分割に属していることを登録する
uniteUF :: UnionFind -> (Int, Int) -> UnionFind
uniteUF uf (i, j)
  | a == b = uf
  | r <= s = IM.insert a b $ IM.insert b (negate $ r + s) uf
  | True   = IM.insert b a $ IM.insert a (negate $ r + s) uf
  where
    (a,r) = getRoot uf i
    (b,s) = getRoot uf j

Submission Info

Compile Error

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

Judge Result