Submission #34030393

Source Code Expand

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import Control.Monad
import qualified Data.ByteString.Char8 as BS
import Data.Char
import Data.List
import Data.Array
import Control.Applicative
main = do
  [n,m,e] <- bsGetLnInts
  uvs <- replicateM e ((\(u:v:_) -> (u,v)) <$> 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] -> [Int]
abc264e n m e uvs q xs = scanr (+) ansQ ccds
 
 
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

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

import Data.Array
import Control.Applicative

main = do
  [n,m,e] <- bsGetLnInts
  uvs <- replicateM e ((\(u:v:_) -> (u,v)) <$> 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] -> [Int]
abc264e n m e uvs q xs = scanr (+) ansQ ccds
  where
-- 先頭 N 個は都市、後半 M 個が発電所
    uf0 = newUF ufF (n+m) (replicate n (1,False) ++ replicate m (0, True))
-- 電線の配列
    uvA = listArray (1,e) uvs
-- 電線が切れないかチェック表
    uvX = accumArray (&&) True (1,e) [(x,False) | x <- xs]
-- 切れない電線を配線する
    uf1 = foldl' {-'-} uniteUF uf0 [uv | (uv,True) <- zip uvs (elems uvX)]
-- 最後まで通電している都市の数
    ansQ = sum [cc | (_,(cc, True)) <- getCont uf1]
-- 逆順に Xi を繋いだufを作りつつ、都市数の増加を追跡
    (_,ccds) = mapAccumR step uf1 (tail xs)
    step uf xi
      | a == b     = (uf, 0)
      | pp1 == pp2 = (uf1, 0)
      | pp1        = (uf1, cc2)
      | otherwise  = (uf1, cc1) -- pp2 = True
      where
        uvx@(ux, vx) = uvA ! xi
        (a, _, (cc1,pp1)) = getRoot uf ux
        (b, _, (cc2,pp2)) = getRoot uf vx
        uf1 = uniteUF uf uvx

ufF (cc1,pp1) (cc2,pp2) = (cc1 + cc2, pp1 || pp2)

-- UnionFind

type UnionFind a = (Array Int (Either Int (Int, a)), a -> a -> a)
-- 配列要素は 親id または ランク(正の数),ペイロード のペア

-- 1からNまでのN要素が独立した初期状態を作る
-- ペイロードの統合関数も保持する
newUF :: (a -> a -> a) -> Int -> [a] -> UnionFind a
newUF f n as = (listArray (1, n) [Right (1, a) | a <- as], f)

-- 補助関数 ノードの根まで辿り、代表id, rank, ペイロードを返す
getRoot :: UnionFind a -> Int -> (Int, Int, a)
getRoot uf@(ar, _) i =
  case ar ! i of
    Left k -> getRoot uf k
    Right (r, v) -> (i, r, v)

idof (i, _, _) = i

-- ふたつのノードが同じ分割に属しているか判定する
findUF :: UnionFind a -> Int -> Int -> Bool
findUF uf i j = idof (getRoot uf i) == idof (getRoot uf j)

-- ふたつのノードが同じ分割に属していることを登録する
uniteUF :: UnionFind a -> (Int, Int) -> UnionFind a
uniteUF uf@(ar, f) (i, j)
  | a == b = uf
  | otherwise =
      case compare r s of
        GT -> (ar // [(b, Left a), (a, Right (r, f p q))], f)
        LT -> (ar // [(a, Left b), (b, Right (s, f p q))], f)
        EQ -> (ar // [(a, Left b), (b, Right (succ s, f p q))], f)
  where
    (a, r, p) = getRoot uf i
    (b, s, q) = getRoot uf j

-- 全てのペイロードを、代表idとペアにして返す
getCont :: UnionFind a -> [(Int, a)]
getCont (ar,_) = [(i,p) | (i, Right (_, p)) <- assocs ar]

Submission Info

Compile Error

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

Judge Result