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 .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 (,n-1) $ -1 ; rk <- newArray (,n-1) $ 
return $ U rt rk
find :: UnionFind s -> Int ->ST s Int
find u x = do r <- readArray (root u) x
if r <  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)

もう，手続き型言語で書けよって感じなコードである．

なんか負けたような気がする．

しかし，ランダムアクセスとか必要だから，こんなもん？