Life Goes On

まあまあだけど楽しんでる方です

Could not deduce (MArray (STUArray s) (Int, Int) (ST s)) from the context ()

引き続きモナドと格闘中。

できるようになったこと。

  • STArray や IOArray を動かせるようになった。
  • Disjoint Set が実装できた。→ 9400 分( 6 日と 12 時間 40 分)かかっていたコードが 1 分で返ってくるようになった。ちなみに STArray で 1 分、STUArray で 45 秒くらい。

まだ分からないこと。

  • ST や IO の使い分け。(単なる性能改善という目的からすると ST を使うべきなのかもしれませんが、IO だと runST とか不要で、簡単に書けるように思います)

う〜ん‥

import Control.Monad.ST
import Control.Monad
import Data.Array.MArray
import Data.Array.ST

type DisjointSet s = (STUArray s Int Int, STUArray s Int Int)

main = print $ euler186 records

euler186 :: [[Int]] -> Int
euler186 cs = runST $ do
    djs <- makeDjs (0, 999999)
    ls <- mapM (unionAndCount djs) $ take 2400000 cs
    return $ (+ 1) $ length $ takeWhile (< 990000) ls

records :: [[Int]]
records = call ss
    where call (a0:a1:as) = if a0 == a1 then call as else [a0,a1] : call as

ss :: [Int]
ss = map (fromInteger . s) [1..55] ++ map (`mod` 1000000) (zipWith (+) ss $ drop 31 ss)
    where s k = mod (100003 - 200003*k + 300007*k^3) 1000000

makeDjs :: (Int, Int) -> ST s (DisjointSet s)
makeDjs (u, v) = do
    ps <- newListArray (u, v) [u..v]
    ss <- newListArray (u, v) [1,1..]
    return (ps, ss)

unionAndCount :: DisjointSet s -> [Int] -> ST s Int
unionAndCount djs [a, b] = do
    pa <- djsFind djs a
    pb <- djsFind djs b
    when (pa /= pb) $ djsUnion djs pa pb
    p <- djsFind djs 534287
    readArray (snd djs) p

djsFind :: DisjointSet s -> Int -> ST s Int
djsFind (ps, ss) c = do
    p <- readArray ps c
    when (c /= p) $ writeArray ps c =<< djsFind (ps, ss) p
    readArray ps c

djsUnion :: DisjointSet s -> Int -> Int -> ST s ()
djsUnion (ps, ss) pa pb = do
    sa <- readArray ss pa
    sb <- readArray ss pb
    when (sa > sb) $ writeArray ps pb pa >> writeArray ss pa (sa+sb)
    when (sa <= sb) $ writeArray ps pa pb >> writeArray ss pb (sa+sb)