Problem 186 (UnionFind,DisjointSet)

http://projecteuler.net/index.php?section=problems&id=186
久々のproject euler.
UnionFind,DisjointSetを知っていれば簡単．
問題は，どう関数型言語で実装するか，である(笑)．

{-# LANGUAGE BangPatterns #-}
import Control.Monad.ST (ST,runST)
import Control.Monad (when,liftM)
import Data.Array.ST (STUArray,newArray,readArray,writeArray)

lagFib :: [Integer]
lagFib = map s [1..55] ++ zipWith add lagFib (drop 31 lagFib)
    where s k = (100003 - 200003*k + 300007*k^3) `mod` 1000000
          add x y = (x + y) `mod` 1000000

main :: IO ()
main =print.runST $ do 
         u <- unionFind (10^6)
         let call !c (a:b:cs)
                 | a == b    = call c cs
                 | otherwise = do s <- size u 524287
                                  if s >= 990000 then return c
                                     else union u a b >> call (c+1) cs
         call 0.map fromInteger $ lagFib

data UnionFind s = U {root::STUArray s Int Int,rank::STUArray s Int Int}

unionFind :: Int -> ST s (UnionFind s)
unionFind n = do rt <- newArray (0,n-1) $ -1 ; rk <- newArray (0,n-1) $ 0
                 return $ U rt rk

find :: UnionFind s -> Int ->ST s Int
find u x = do r <- readArray (root u) x
              if r < 0 then return x
                       else do s <- find u r
                               writeArray (root u) x s
                               return s

size :: UnionFind s -> Int -> ST s Int
size u x = find u x >>= liftM negate.readArray (root u)

union :: UnionFind s -> Int -> Int -> ST s ()
union u x y = do 
  px <- find u x              ; py <- find u y
  sx <- readArray (root u) px ; sy <- readArray (root u) py
  rx <- readArray (rank u) px ; ry <- readArray (rank u) py
  case compare rx ry of
    GT -> writeArray (root u) py px >> writeArray (root u) px (sx+sy)
    LT -> writeArray (root u) px py >> writeArray (root u) py (sx+sy)
    EQ -> when (px/=py) $ do writeArray (root u) py px
                             writeArray (root u) px (sx+sy)
                             writeArray (rank u) px (rx+1)

もう，手続き型言語で書けよって感じなコードである．
なんか負けたような気がする．
しかし，ランダムアクセスとか必要だから，こんなもん？