Union Find in Haskell, using the State Monad.

UnionFind.hs

import Control.Monad.Trans.State.Strict
import Control.Monad
import Data.Ord
import Data.Maybe
import qualified Data.IntMap.Strict as M

data S = S { parentM :: M.IntMap Int
           , rankM   :: M.IntMap Int }

modifyP f = modify $ \s -> s { parentM = f (parentM s) }
modifyR f = modify $ \s -> s { rankM   = f (rankM s)   }

parent :: Int -> State S Int
parent k = do
  m <- parentM <$> get
  return $ maybe k id (M.lookup k m)

rank :: Int -> State S Int
rank k = do
  m <- rankM <$> get
  return $ maybe 0 id (M.lookup k m)

find :: Int -> State S Int
find x = do
  p <- parent x
  when (p /= x) $ do
    p' <- find p
    modifyP $ M.insert x p'
  parent x

union :: Int -> Int -> State S ()
union x y = do
  m <- get
  xroot <- find x
  yroot <- find y
  xrank <- rank xroot
  yrank <- rank yroot
  case xrank `compare` yrank of
    LT -> modifyP $ M.insert xroot yroot
    GT -> modifyP $ M.insert yroot xroot
    EQ -> do modifyP $ M.insert yroot xroot
             modifyR $ M.adjust (+1) xroot