■

先日UTPC2013に参加した。
東京大学プログラミングコンテスト2013

問題CでTLEにはまったのでメモ。

問題C概要

２つの連結グラフG_1, G_2が与えられる。
一つの辺を使ってその２つのグラフを連結するときに、連結したグラフの直径の最大値、最小値を求めよ。

制約

それぞれのグラフの頂点数<= 10^3, 辺数<= 10^4

解法

G_1に対して各頂点vから最も遠い頂点までの距離をd_1[v]とする。G_2に対しても同様にd_2[v]とする。
min(d_1) = r1　min(d_2) = r2
max(d_1) = R1　max(d_2) = R2
とする。
連結したグラフの直径の最大値はR1+R2+1、
最小値はmax(R1,R2,r1+r2+1)となる。

実装

d_1,d_2を求めるのには幅優先探索をすればいいのだが
次のようなコードを書いたらTLEした。
http://utpc2013.contest.atcoder.jp/submissions/137893

dist :: Int -> [(Int,Int)] -> UArray Int Int
dist n es = listArray (0,n-1) [ f v | v <- [0..n-1] ]
    where
    g :: Array Int [Int]
    g = accumArray (flip (:)) [] (0,n-1) $ es ++ map swap es
    f v = runST $ do
        r <- newArray (0,n-1) False
        go r [v] [] 0 0
    go :: STUArray s Int Bool -> [Int] ->  [Int] -> Int -> Int -> ST s Int
    go _ [] [] d !acc = return acc
    go vis [] ns d !acc = go vis ns [] (d+1) acc
    go vis (v:l) !ns !d !acc = do
        b <- readArray vis v
        if b then
            go vis l ns d acc
        else do
            writeArray vis v True
            let acc' = max acc d
                ns' = g ! v ++ ns
            go vis l ns' d acc'

キューから一つ頂点を取り出し、
その頂点をそれまでに訪れたことがあるかどうかを調べ、
それまでに訪れていなければその頂点に隣接している頂点をキューに追加する。

このコードはすべての辺に対して両方の頂点を一回ずつキューに追加するので
効率が悪い。（高々２倍程度なので通常のコンテストならこのコードでも通るはずだが）

キューから取り出した頂点vに隣接した各頂点uに対して
uをまだ訪れていないならばuをキューに追加する。
というアルゴリズムにすると、キューに追加される頂点数が
グラフの頂点数と等しくなり高速する。
http://utpc2013.contest.atcoder.jp/submissions/140760

dist :: Int -> [(Int,Int)] -> [Int]
dist n es = do
    v <- [0..n-1]
    return $ runST $ do
        r <- newArray (0,n-1) True :: ST s (STUArray s Int Bool)
        unsafeWrite r v False
        go r 1 [v] []
    where
    g :: Array Int [Int]
    g = accumArray (flip (:)) [] (0,n-1) $ es ++ map swap es
    go r d [] [] = return (d-1)
    go r d [] !ns = go r (d+1) ns []
    go r d (v:vs) !ns = gosub (unsafeAt g v) ns
        where
        gosub [] !ns = go r d vs ns
        gosub (u:us) !ns = do
            b <- unsafeRead r u
            if b then do
                unsafeWrite r u False
                gosub us (u:ns)
            else
                gosub us ns

同様のアルゴリズムをjavaで実装したところ同程度の実行時間になった。
http://utpc2013.contest.atcoder.jp/submissions/140770

余談

他の人のHaskellコードを見ていたらmkotha氏とcojna氏が400msで動く
非常に速いコードを書いていた。
mkotha氏: http://utpc2013.contest.atcoder.jp/submissions/140275
cojna氏: http://utpc2013.contest.atcoder.jp/submissions/140394
見た感じキューを自作していたり、Unboxed Tuple を使っていたりしていて怖い。