We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
  1. Prepare
  2. Functional Programming
  3. Functional Structures
  4. Prison Transport
  5. Discussions

Prison Transport

  • mykhail0
     + 0 comments

    This is a Disjoint Set problem, one of the few data structures that are specifically imperative, so I wonder why this problem is in the "Functional Structures" category. Anyways, here is my Haskell solution using mutable arrays to implement DSU:

    import Control.Monad (replicateM_)
import Data.List (sort, group)
import Data.Array.IO

main :: IO ()
main = do
  g <- (initDSU =<< readLn)
  m <- readLn
  replicateM_ m $ addEdge g
  getSizes g >>= print . cost

addEdge :: DSU -> IO ()
addEdge g = do
  pq <- getLine
  let [p, q] = fmap read $ words pq
  unionSets p q g

-- Given: sizes of connected components.
cost :: [Int] -> Int
cost = sum . fmap (ceiling . sqrt . fromIntegral)

data DSU = DSU {
  _n :: Int,
  parent :: IOUArray Int Int,
  size :: IOUArray Int Int
}

setSize :: DSU -> Int -> Int -> IO ()
setSize = writeArray . size

setParent :: DSU -> Int -> Int -> IO ()
setParent = writeArray . parent

getSize :: DSU -> Int -> IO Int
getSize = readArray . size

getParent :: DSU -> Int -> IO Int
getParent = readArray . parent

initDSU :: Int -> IO DSU
initDSU n = do
  p <- mkArray [1..n]
  s <- mkArray $ replicate n 1
  pure $ DSU n p s
  where mkArray = newListArray (1, n)

findSet :: Int -> DSU -> IO Int
findSet v g = do
  p <- getParent g v
  if v == p then pure p
  else do
    r <- findSet p g
    setParent g v r
    pure r

unionSets :: Int -> Int -> DSU -> IO ()
unionSets a b g = do
  a' <- findSet a g
  b' <- findSet b g
  if a' == b' then pure ()
  else do
    sa <- getSize g a'
    sb <- getSize g b'
    let (ar, br) = if sa < sb then (b', a') else (a', b')
    setParent g br ar
    setSize g ar $ sa + sb

getReprs :: DSU -> IO [Int]
getReprs g@(DSU n _ _) = mapM (flip findSet g) [1..n] >>=
  pure . fmap head . group . sort

getSizes :: DSU -> IO [Int]
getSizes g = getReprs g >>= mapM (getSize g)