AtCoder Typical Contest 001

Submission #420313

Source codeソースコード

{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NumDecimals #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExplicitForAll #-}

import Control.Applicative
import Control.Monad
import qualified Control.Monad.Primitive as Primitive
import qualified Data.ByteString.Char8 as BS
import Data.List
import qualified Data.Vector.Unboxed.Mutable as UM
import GHC.Exts (Int(..))

main :: IO ()
main = do
  [n, q] <- getInts
  uf <- newUf n
  replicateM_ q $ do
    [p, a, b] <- getInts
    case p of
      0 -> joinUf uf a b
      _ -> do
        r <- sameUf uf a b
        putStrLn $ case r of
          True -> "Yes"
          False -> "No"

----------------------------------------------------------------------------
-- IO

getInts :: IO [Int]
getInts = readInts <$> BS.getLine

readInts :: BS.ByteString -> [Int]
readInts = unfoldr $ \(BS.dropWhile (==' ') -> s) ->
  if s == "" then Nothing else case BS.readInt s of
    Just z -> Just z
    _ -> error $ "not an integer: " ++ show s

----------------------------------------------------------------------------
-- UnionFindST

newtype UnionFind s = UnionFind (UM.MVector s Int)

sameUf
  :: (Primitive.PrimMonad m, Applicative m)
  => UnionFind (Primitive.PrimState m)
  -> Int
  -> Int
  -> m Bool
sameUf uf x y = (==) <$> lookupUf uf x <*> lookupUf uf y
{-# INLINE sameUf #-}

lookupUf
  :: (Primitive.PrimMonad m, Applicative m)
  => UnionFind (Primitive.PrimState m)
  -> Int
  -> m Int
lookupUf !(UnionFind v) idx_
  = checkBoundsUf "lookupUf" v idx_ $
    lookupUf' v idx_
{-# INLINE lookupUf #-}

joinUf
  :: (Primitive.PrimMonad m, Applicative m)
  => UnionFind (Primitive.PrimState m)
  -> Int
  -> Int
  -> m ()
joinUf (UnionFind v) x y =
  checkBoundsUf "joinUf[x]" v x $
  checkBoundsUf "joinUf[y]" v y $ do
  rx <- lookupUf' v x
  ry <- lookupUf' v y
  when (rx /= ry) $ do
    rankx <- UM.unsafeRead v rx
    ranky <- UM.unsafeRead v ry
    case compare rankx ranky of
      GT -> UM.unsafeWrite v rx ry
      LT -> UM.unsafeWrite v ry rx
      EQ -> do
        UM.unsafeWrite v rx ry
        UM.unsafeWrite v ry (rankx - 1)
{-# INLINE joinUf #-}

checkBoundsUf :: String -> UM.STVector s Int -> Int -> a -> a
checkBoundsUf loc vec idx_ body
  | idx_ < 0 || idx_ >= UM.length vec = error $
    loc ++ ": index out of bounds(size=" ++ show (UM.length vec) ++ "): " ++ show idx_
  | otherwise = body
{-# INLINE checkBoundsUf #-}

lookupUf'
  :: forall m
  .  (Primitive.PrimMonad m, Applicative m)
  => UM.STVector (Primitive.PrimState m) Int
  -> Int
  -> m Int
lookupUf' v i0 = loop i0
  where
    loop :: Int -> m Int
    loop !i = Primitive.primitive $ \s -> case loop' i s of
      (# s', r #) -> (# s', I# r #)
    {-# INLINE loop #-}

    loop' !i s = case Primitive.internal (loop'' i) s of
      (# s', I# r #) -> (# s', r #)

    loop'' !i = do
      r <- UM.unsafeRead v i
      if r < 0
        then return i
        else do
          r' <- loop r
          UM.unsafeWrite v i r'
          return r'
    {-# INLINE loop'' #-}
{-# INLINE lookupUf' #-}

newUf
  :: (Primitive.PrimMonad m, Applicative m)
  => Int
  -> m (UnionFind (Primitive.PrimState m))
newUf !size
  | size < 0 = error $ "newUf: negative size: " ++ show size
  | otherwise = UnionFind <$> UM.replicate size (-1)
{-# INLINE newUf #-}

Submission

Task問題 B - Union Find
User nameユーザ名 mkotha
Created time投稿日時
Language言語 Haskell (Haskell Platform 2014.2.0.0)
Status状態 AC
Score得点 100
Source lengthソースコード長 3773 Byte
File nameファイル名
Exec time実行時間 289 ms
Memory usageメモリ使用量 2804 KB

Test case

Set

Set name Score得点 / Max score Cases
Sample - 00_sample_01.txt
All 100 / 100 00_sample_01.txt,subtask_01_01.txt,subtask_01_02.txt,subtask_01_03.txt,subtask_01_04.txt,subtask_01_05.txt,subtask_01_06.txt,subtask_01_07.txt,subtask_01_08.txt,subtask_01_09.txt,subtask_01_10.txt,subtask_01_11.txt,subtask_01_12.txt,subtask_01_13.txt,subtask_01_14.txt,subtask_01_15.txt,subtask_01_16.txt,subtask_01_17.txt,subtask_01_18.txt

Test case

Case name Status状態 Exec time実行時間 Memory usageメモリ使用量
00_sample_01.txt AC 29 ms 1308 KB
subtask_01_01.txt AC 178 ms 1944 KB
subtask_01_02.txt AC 29 ms 2204 KB
subtask_01_03.txt AC 264 ms 1904 KB
subtask_01_04.txt AC 289 ms 2716 KB
subtask_01_05.txt AC 42 ms 1820 KB
subtask_01_06.txt AC 43 ms 2648 KB
subtask_01_07.txt AC 269 ms 1944 KB
subtask_01_08.txt AC 283 ms 2704 KB
subtask_01_09.txt AC 29 ms 1432 KB
subtask_01_10.txt AC 29 ms 2260 KB
subtask_01_11.txt AC 261 ms 1912 KB
subtask_01_12.txt AC 286 ms 2804 KB
subtask_01_13.txt AC 223 ms 1940 KB
subtask_01_14.txt AC 31 ms 2688 KB
subtask_01_15.txt AC 265 ms 1912 KB
subtask_01_16.txt AC 288 ms 2712 KB
subtask_01_17.txt AC 222 ms 2712 KB
subtask_01_18.txt AC 281 ms 2708 KB