Submission #420313

---

Source Code Expand

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

Judge Result