Submission #420313


Source Code Expand

Copy
{-# 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 Info

Submission Time
Task B - Union Find
User mkotha
Language Haskell (Haskell Platform 2014.2.0.0)
Score 100
Code Size 3773 Byte
Status
Exec Time 289 ms
Memory 2804 KB

Judge Result

Set Name Score / Max Score Test Cases
Sample 0 / 0 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
Case Name Status Exec Time Memory
00_sample_01.txt 29 ms 1308 KB
subtask_01_01.txt 178 ms 1944 KB
subtask_01_02.txt 29 ms 2204 KB
subtask_01_03.txt 264 ms 1904 KB
subtask_01_04.txt 289 ms 2716 KB
subtask_01_05.txt 42 ms 1820 KB
subtask_01_06.txt 43 ms 2648 KB
subtask_01_07.txt 269 ms 1944 KB
subtask_01_08.txt 283 ms 2704 KB
subtask_01_09.txt 29 ms 1432 KB
subtask_01_10.txt 29 ms 2260 KB
subtask_01_11.txt 261 ms 1912 KB
subtask_01_12.txt 286 ms 2804 KB
subtask_01_13.txt 223 ms 1940 KB
subtask_01_14.txt 31 ms 2688 KB
subtask_01_15.txt 265 ms 1912 KB
subtask_01_16.txt 288 ms 2712 KB
subtask_01_17.txt 222 ms 2712 KB
subtask_01_18.txt 281 ms 2708 KB