Kattis solution is not fast enough (MArray) [Solved]

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)

Nothing immediately jumps out at me, but if there’s a huge number of queries being sent, you might actually be hitting timeouts due to String being too slow. I’ve definitely hit that before when using Haskell for problems on other similar platforms. Another bit of low hanging fruit - Query is probably lazier than you’d like. You might want

data Query = Merge !Int !Int | Size !Int
1 Like

You were right! String was at fault! I didn’t expect that since I just pattern match and traverse the strings (and String is nice for that). I guess String is just inherently inefficient then.

The solution passes the tests now! Here is the final version, for anyone curious:

{-# LANGUAGE OverloadedStrings #-}

module Main where

import Data.Array.IO (IOUArray)
import Data.Array.MArray
import Data.Foldable (traverse_)
import Data.Text hiding (head, map, tail)
import Data.Text.IO
import Data.Text.Read
import Prelude hiding (getContents, lines, read, words)

main :: IO ()
main = do
  wordRows <- map words . lines <$> getContents
  uf <- newUnionFind (decimalU (head (head wordRows)) + 1)
  traverse_ (act uf . mkQuery) (tail wordRows)

data Query = Merge !Int !Int | Size !Int deriving (Show)

mkQuery :: [Text] -> Query
mkQuery ["t", a, b] = Merge (decimalU a) (decimalU b)
mkQuery ["s", a] = Size (decimalU a)

decimalU :: Text -> Int
decimalU = fst . either error id . decimal

...

Yes, I was also going to suggest String is the culprit. Strings are linked lists of characters which incurs a lot of overhead; any time the input or output is as big as 10^5 or so (this problem, with an input of size 10^6, has particularly massive I/O), you do not want to be using String for parsing the input. Personally, since inputs are pretty much always guaranteed to be ASCII, I always reach for Data.ByteString.Lazy.Char8 rather than Data.Text, though I don’t know how much of a difference it makes.

3 Likes

Hi Brent! Super cool to have you chime in. I’m actually following your list of solved problems to have an idea of which to do!

I appreciate the tip - I’ll look into using Char8 instead.

Hope you and your programming team are looking forward to Advent of Code as much as I am!

1 Like

You know, I’ve been so heads-down in the end of the semester that I hadn’t even thought about Advent of Code, but now that you mention it, yes, I’m very much looking forward to it!

1 Like