LabmdaCase は Golf の味方

競プロ典型 90 問 12 日目 - Red Painting（★4）

{-# LANGUAGE LambdaCase #-}

import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Bool
import qualified Data.ByteString.Char8 as C
import qualified Data.IntSet as S
import Data.List

main = do
  [h,w] <- get
  q <- readLn
  qs <- replicateM q get
  mapM_ (putStrLn  . bool "No" "Yes") $ sol (h+2) (w+2) qs

get = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine

sol h w qs = runST $ do
  uf <- newUF (h*w)
  let
    f (cs,ps) = \case
      (1:r:c:_) -> do
        let p = r*w+c
        forM_ [1,-1,w,-w] $ \d -> when ((p+d) `S.member` ps) $ unite uf p (p+d)
        return (cs,S.insert p ps)
      (2:ra:ca:rb:cb:_) -> do
        let
          a = ra*w+ca
          b = rb*w+cb
        c <- clade uf a b
        return ((a `S.member` ps && b `S.member` ps && c):cs,ps)
  reverse . fst <$> foldM f ([],S.empty) qs

data UnionFind s = UnionFind {ids:: STUArray s Int Int, szs:: STUArray s Int Int}

newUF :: Int -> ST s (UnionFind s)
newUF n = liftM2 UnionFind (newListArray (1,n) [1..n]) (newArray (1,n) 1)

clade uf p q = (==) <$> root uf p <*> root uf q

root :: UnionFind s -> Int -> ST s Int
root uf i = do
  ix <- readArray (ids uf) i
  if ix == i
    then return i
    else do
      readArray (ids uf) ix >>= writeArray (ids uf) i
      root uf ix

unite :: UnionFind s -> Int -> Int -> ST s ()
unite uf p q = do
  i <- root uf p
  j <- root uf q
  unless (i==j) $ do
    szi <- readArray (szs uf) i
    szj <- readArray (szs uf) j
    if szi<szj
      then do
        writeArray (ids uf) i j
        writeArray (szs uf) j (szi+szj)
      else do
        writeArray (ids uf) j i
        writeArray (szs uf) i (szj+szi)