Hey! I’m still trying to solve puzzles on Kattis, but this one got me confused. The problem: tildes.
Anyone will recognize that it’s about the union-find data structure, so it was the first time I tried to use it in Haskell. Inspiration for the algorithm from Kwang’s blog post. I just made it use MArray
instead, and added a condition in the unite
block. Also, the find
function wasn’t used at all, so I removed it.
What’s weird is that it’s not performing well enough to pass the time limit, and I don’t see what the culprit could be. Any guesses?
module Main where
import Data.Array.IO (IOUArray)
import Data.Array.MArray
import Data.Foldable (traverse_)
main :: IO ()
main = do
(n : _) : rest <- map words . lines <$> getContents
uf <- newUnionFind (read n + 1)
traverse_ (act uf . mkQuery) rest
data Query = Merge Int Int | Size Int deriving (Show)
mkQuery :: [String] -> Query
mkQuery ["t", a, b] = Merge (read a) (read b)
mkQuery ["s", a] = Size (read a)
act :: UnionFind -> Query -> IO ()
act uf (Merge a b) = merge uf a b
act uf (Size a) = do
i <- root uf a
sz <- readArray (szs uf) i
print sz
data UnionFind = UnionFind
{ ids :: IOUArray Int Int,
szs :: IOUArray Int Int
}
newUnionFind :: Int -> IO UnionFind
newUnionFind n = UnionFind <$> newListArray (0, n - 1) [0 .. n - 1] <*> newArray (0, n - 1) 1
root :: UnionFind -> Int -> IO Int
root uf i = do
id <- readArray (ids uf) i
if id /= i
then do
gpid <- readArray (ids uf) id
writeArray (ids uf) i gpid
root uf id
else return i
merge :: UnionFind -> Int -> Int -> IO ()
merge uf p q = do
i <- root uf p
j <- root uf q
szi <- readArray (szs uf) i
szj <- readArray (szs uf) j
if i == j
then return ()
else
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)