29 Nov 2018, 22:06

Haskellで蟻本(初級編) - GCJの問題に挑戦してみよう(1)

※蟻本の入力例でしかテストしていません

下2つはコンテストで出たとして解けなそう…。

Minimum Scalar Product

{-# LANGUAGE BangPatterns, FlexibleContexts #-}

import Control.Monad
import Control.Monad.ST
import Data.Array
import Data.Array.ST
import Data.STRef
import Data.List

-- Minimum Scalar Product (2008 Round1A A)
q1 :: Int -> [Int] -> [Int] -> Int
q1 n xs ys = sum $ zipWith (*) xs' ys'
  where
    xs' = sort xs
    ys' = sortBy (flip compare) ys
  • 蟻本の解説参照

Crazy Rows

-- Crazy Rows (2009 Round2 A)
q2 :: Int -> [String] -> Int
q2 n xss = runST $ do
    a <- newListArray (1, n) (fmap last1 xss) :: ST s (STUArray s Int Int)
    ans <- newSTRef 0
    forM_ [1 .. n] $ \i -> do
        ai <- readArray a i
        when (ai > i) $ do
            js <- flip filterM [i + 1 .. n] $ \j -> do
                aj <- readArray a j
                return $ aj <= i
            when (not $ null js) (swap ans a i (head js))
    readSTRef ans
  where
    last1 xs = length $ dropWhile (/= '1') (reverse xs)
    swap ans a i j
        | i == j    = return ()
        | otherwise = do
            aj  <- readArray a j
            writeArray a j =<< readArray a (j - 1)
            writeArray a (j - 1) aj
            modifySTRef' ans (+ 1)
            swap ans a i (j - 1)
  • MArray面倒くさい

Bribe the Prisoners

-- Bribe the Prisoners (2009 Round 1C C)
q3 :: Int -> Int -> [Int] -> Int
q3 p q as = runST $ do
    dp <- newArray ((0, 0), (q + 1, (q + 1))) maxBound :: ST s (STUArray s (Int, Int) Int)
    sequence_ [writeArray dp (i, i + 1) 0 | i <- [0 .. q]]
    forM_ [2 .. q + 1] $ \w -> do
        forM_ [0 .. (q + 1) - w] $ \i -> do
            let j = i + w
            t <- fmap minimum $ forM [i + 1 .. j - 1] $ \k -> do
                dpik <- readArray dp (i, k)
                dpkj <- readArray dp (k, j)
                return (dpik + dpkj)
            writeArray dp (i, j) (t + a ! j - a ! i - 1 - 1)
    readArray dp (0, q + 1)
  where
    a = listArray (0, q + 1) (0 : as ++ [p + 1])
  • dp[i][j] := (i, j)を解放するのに必要な金貨
    • 例2) P = 20, Q = 3, A = {3, 6, 14}
      • 配列Aを作る時に両端を追加してA = [0, 3, 6, 14, 21]
      • 例えばdp[0][3]は左端と囚人14は解放済みとして、囚人3と囚人6を解放するのに必要な金貨の最小枚数
  • まず幅が1の時を0で初期化する
    • 幅が1、例えばdp[2][3]は間に解放すべき囚人がいないため金貨不要
  • 幅が2(間に1人解放すべき囚人が存在する)から幅Q+1までループ
  • 最初に解放する囚人をすべて試し、最小コストのものを探す
    • 例えばdp[0][3]dp[0][1] + dp[1][3]dp[0][2] + dp[2][3]の小さい方
    • 先に囚人3と囚人6のどちらを解放するにしても14 - 0 - 1 - 1の金貨は必要

Millionaire

-- Millionaire (2008 APAC local onsites C)
q4 :: Int -> Double -> Int -> Double
q4 m p x = go m gx
  where
    m2 = 2 ^ m
    go 0 j
        | j == m2   = 1.0
        | otherwise = 0.0
    go i j = maximum $ do
        v <- [0 .. min j (m2 - j)]
        return $ p * (memo ! (i - 1, j + v)) + (1 - p) * (memo ! (i - 1, j - v))
    memo = listArray ((0, 0), (m, m2)) [go i j | i <- [0 .. m], j <- [0 .. m2]]
    gx = (x * m2) `div` 1000000
  • 例としてM = 2, P = 0.5, X = 500000を考える
  • dp[i][j] := 残りラウンドがiで、所持金がグループjのとき、最善の戦略をとってお金を持って帰れる確率
  • 所持金のグループ
    • ラウンド数が2のとき
      • j = 0, 0 ~ 249,999
      • j = 1, 250,000 ~ 499,999
      • j = 2, 500,000 ~ 749,999
      • j = 3, 750,000 ~ 999,999
      • j = 4, 1,000,000 ~
  • 配列を埋め終わった状態はこんな感じになる

            [0]    [1]    [2]    [3]    [4]
    dp[0] | 0.00 | 0.00 | 0.00 | 0.00 | 1.00 |
    dp[1] | 0.00 | 0.00 | 0.50 | 0.50 | 1.00 |
    dp[2] | 0.00 | 0.25 | 0.50 | 0.75 | 1.00 |
    
  • 例えばdp[1][2]を更新するときはどうなるか

    • イメージ的にはdp[1][2]から見てV字の先の値が必要になって、掛けた金額に応じてV字の角度が変化する(意味不明)
    • 掛けた金額が0なら、勝とうが負けようがdp[0][2]の状態になって
    • 全額かけたなら、勝てばdp[0][4]の状態になって負ければdp[0][0]の状態になって
    • その間の金額をかけたら、dp[0][1], dp[0][3]みたいな…
  • 最初に持っている金額500,000はグループ2なので答えはdp[2][2]0.50

    • 蟻本のint i = (ll)X * n / 1000000;の部分で最初の所持金がどのグループか判別している
    • 500,000 / (1,000,000 / 2^2) == 500,000 * 2^2 / 1,000,000 == 2

27 Nov 2018, 22:42

Haskellで蟻本(初級編) - 数学的な問題を解くコツ

参考資料


※蟻本の入力例でしかテストしていません

ユークリッド互除法

{-# LANGUAGE BangPatterns #-}

import qualified Data.IntMap.Strict as IM

-- ユークリッド互除法
gcd' :: Int -> Int -> Int
gcd' x y = go x y
  where
    go a 0 = a
    go a b = go b (a `mod` b)

-- 線分上の格子点の個数
q1 :: (Int, Int) -> (Int, Int) -> Int
q1 (x1, y1) (x2, y2) = gcd' (abs $ x1 - x2) (abs $ y1 - y2) - 1
  • 参考資料の「1.10.2.1 ユークリッドの互除法」を読む
  • 証明も理解する

拡張ユークリッド互除法

-- 拡張ユークリッド互除法
extgcd :: Int -> Int -> (Int, Int)
extgcd x y = go x y
  where
    go a 0 = (1, 0)
    go a b = (t, s - q * t)
      where
        (q, r) = a `divMod` b
        (s, t) = extgcd b r

-- 双六
q2 :: Int -> Int -> Maybe (Int, Int, Int, Int)
q2 a b
    | gcd' a b == 1 = case () of
        _ | x < 0, y < 0 -> Just (0, 0, abs x, abs y)
          | x < 0        -> Just (0, y, abs x, 0)
          | y < 0        -> Just (x, 0, 0, abs y)
          | otherwise    -> Just (x, y, 0, 0)
    | otherwise     = Nothing
  where
    (x, y) = extgcd a b
  • 拡張ユークリッド互除法で何ができるの?

    • ax + by = gcd(a, b)の整数解、(x, y)を求めることができる
  • 上記のプログラムが理解できないんだけど?(蟻本の解説本が欲しいんだけど?)

    参考資料の「1.10.3 整数論の基本定理」を読む
    整数a,b が互いに素のとき ax + by = 1 は整数解をもつ
    整数a,b が互いに素のとき ax + by = c は整数解をもつ
    整数a, b が互いに素でなくても c が gcd(a, b) で割り切れるなら ax + by = c は整数解をもつ
    ということはax + by = gcd(a, b) は整数解をもつ
    a を b で割って a = bq + r として上に式に代入すると
    b(qx + y) + rx = gcd(a, b)
    s = qx + y, x = t とおいて bs + rt = gcd(s, t) = gcd(a, b)
    (b, r) は (a, b) よりも小さいため再帰的に解くことができる
    (s, t) がわかっているなら上の式で元の解(x, y) がわかる
    停止条件は r == 0 のとき bs + 0t = gcd(b, 0) の解(s, t) = (1, 0)
    

素数に関する基本的なアルゴリズム

-- 素数判定
isPrime :: Int -> Bool
isPrime n = all ((/= 0) . (n `mod`)) $ takeWhile ((<= n) . (^ 2)) (primes n)

-- n以下の素数
primes :: Int -> [Int]
primes n = 2 : 3 : filter isPrime xs
  where
    xs = takeWhile (<= n) [6 * i + j | i <- [1 ..], j <- [-1, 1]]

-- 約数の列挙
divisor :: Int -> [Int]
divisor n = foldr f [] $ takeWhile ((<= n) . (^ 2)) [1 .. n]
  where
    f x ds
        | r == 0, q /= x = x : q : ds
        | r == 0         = x : ds
        | otherwise      = ds
      where
        (q, r) = n `divMod` x

-- 素因数分解
primeFactor :: Int -> IM.IntMap Int
primeFactor n = go IM.empty n $ takeWhile ((<= n) . (^ 2)) (primes n)
  where
    go im 1  _  = im
    go im !k [] = IM.insertWith (+) k 1 im
    go im !k (p:ps)
        | r == 0    = go (IM.insertWith (+) p 1 im) q (p:ps)
        | otherwise = go im k ps
      where
        (q, r) = k `divMod` p

エラトステネスの篩

-- エラトステネスの篩
sieve :: Int -> [Int]
sieve n = takeWhile (<= n) (sieve' [2 ..])
  where
    sieve' (x:xs) = x : sieve' (filter ((/= 0) . (`mod` x)) xs)

-- 区間内の素数の個数
q5 :: Int -> Int -> Int
q5 a b = go [a .. b - 1] (takeWhile ((< b) . (^ 2)) [2 ..])
  where
    go !ps []     = length ps
    go !ps (x:xs) = go (f ps) (f xs)
      where
        f as = filter ((/= 0) . (`mod` x)) as
  • Haskellっぽくてかっこいいけど遅い、多分TLE

べき乗を高速に計算する

-- べき乗を高速に計算する
modPow' :: Int -> Int -> Int -> Int
modPow' x n m = go 1 x n
  where
    go !k !_ 0 = k
    go !k !p q
        | r == 1    = go (k * p `mod` m) (p ^ 2 `mod` m) q'
        | otherwise = go k               (p ^ 2 `mod` m) q'
      where
        (q', r) = q `divMod` 2

modPow :: Int -> Int -> Int -> Int
modPow x n m = go n
  where
    go 0 = 1
    go k
        | odd k     = (t * x) `mod` m
        | otherwise = t
      where
        s = go (k `div` 2)
        t = (s * s) `mod` m

-- Carmichael Numbers (UVa No.10006)
q6 :: Int -> Bool
q6 n
    | isPrime n = False
    | otherwise = all f [2 .. n - 1]
  where
    f x = modPow x n n == x `mod` n
  • modPowのほうが分かりやすくて好き

23 Nov 2018, 17:19

Haskellで蟻本(初級編) - グラフ

※蟻本の入力例でしかテストしていません

TODO: コードをもっと綺麗にする

入力の形式について

標準入力は下記の形式で与えられるものとする。

  • n: 頂点数
  • m: 辺の数
  • c: 重み

    n m
    f1 t1
    f2 t2
    ...
    fn tn
    
    n m
    f1 t1 c1
    f2 t2 c2
    ...
    fn tn cn
    

グラフの表現

隣接行列も隣接リストもaccumArrayを使えば良い。

隣接行列

{-# LANGUAGE BangPatterns, FlexibleContexts #-}

import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Array
import Data.STRef
import Data.List

-- 隣接行列
buildGMatrix :: (Int, Int) -> [(Int, Int)] -> Array (Int, Int) Int
buildGMatrix (i, j) es = accumArray (flip const) 0 ((i, i), (j, j)) (zip es (repeat 1))

隣接リスト

-- 隣接リスト
buildGList :: (Int, Int) -> [(Int, a)] -> Array Int [a]
buildGList (i, j) es = accumArray (flip (:)) [] (i, j) es
  • containersData.Graph.buildGと同じ

グラフの探索

-- 二部グラフ判定
q1 :: Int -> Int -> [(Int, Int)] -> Bool
q1 n m es = runST $ do
    color <- newArray (0, n - 1) 0 :: ST s (STUArray s Int Int)
    bs <- forM [0 .. n - 1] $ \i -> do
        ci <- readArray color i
        if ci == 0
            then go color 1 i
            else return True
    return (and bs)
  where
    g = buildGList (0, n - 1) es
    go color c i = do
        ci <- readArray color i
        if ci == 0
            then do
                writeArray color i c
                and <$> sequence (fmap (go color (-c)) (g ! i))
            else return (ci == c)

最短経路問題

ベルマンフォード法

-- ベルマン-フォード法
bellman :: Int -> Int -> Int -> [(Int, Int, Int)] -> Array Int Int
bellman n m start es = runSTArray $ do
    d <- newArray (0, n - 1) maxBound :: ST s (STArray s Int Int)
    writeArray d start 0
    go d
    return d
 where
    go d = do
        rupd <- newSTRef False
        forM_ es $ \(f, t, c) -> do
            df <- readArray d f
            dt <- readArray d t
            when (df /= maxBound && df + c < dt) $ do
                writeArray d t (df + c)
                writeSTRef rupd True
        readSTRef rupd >>= \upd -> when upd (go d)
  • 用意するもの
    • 辺のリスト
    • 1次元配列
      • 初期化:始点のみ0, 残りはINF
      • maxBoundに加算すると負になってしまうので注意
  • ループ
    • 辺のリストに対してループ
    • 辺の始点と終点の距離を配列から取得し、終点の距離を更新
    • 1つでも更新できたら再帰呼び出し

ダイクストラ法

-- ダイクストラ法
dijkstra :: Int -> Int -> Int -> [(Int, Int, Int)] -> Array Int Int
dijkstra n m start es = runSTArray $ do
    d <- newArray (0, n - 1) maxBound :: ST s (STArray s Int Int)
    writeArray d start 0
    go d (push (0, start) Empty :: Skew (Int, Int))
    return d
  where
    g = buildGList (0, n - 1) [(a, (b, c)) | (a, b, c) <- es]
    go d que
        | Just ((c, f), q) <- pop que = do
            rq <- newSTRef q
            df <- readArray d f
            when (c <= df) $ do
                forM_ (g ! f) $ \(t, c) -> do
                    dt <- readArray d t
                    when (df + c < dt) $ do
                        writeArray d t (df + c)
                        modifySTRef' rq (\q -> push (df + c, t) q)
            go d =<< readSTRef rq
        | otherwise = return ()
  • 用意するもの
    • グラフ(隣接リスト)
    • 優先度付きキュー
      • (コスト, 頂点)の並びでタプルにし、コストの小さい順で取り出せるようにする
      • 初期化:(0, 始点)
    • 1次元配列
      • 初期化:始点のみ0, 残りはINF
  • ループ
    • キューが空になるまでループ
    • キューのコストと配列のコストを比較し配列のほうが既に小さいときはスキップ
    • グラフより隣接する頂点を取り出し、更新できるときは
      • 配列を更新
      • キューにも追加

ワーシャル-フロイド法

-- ワーシャル-フロイド法
warshall :: Int -> Int -> [(Int, Int, Int)] -> Array (Int, Int) Int
warshall n m es = runSTArray $ do
    d <- newArray ((0, 0), (n - 1, n - 1)) inf
    sequence_ [writeArray d (i, i) 0 | i <- [0 .. n - 1]]
    sequence_ [writeArray d (f, t) c | (f, t, c) <- es]
    forM_ [0 .. n - 1] $ \k -> do
        forM_ [0 .. n - 1] $ \i -> do
            forM_ [0 .. n - 1] $ \j -> do
                dij <- readArray d (i, j)
                dik <- readArray d (i, k)
                dkj <- readArray d (k, j)
                writeArray d (i, j) $ min dij (dik + dkj)
    return d
  where
    inf = 10 ^ 9 + 7
  • 用意するもの
    • 辺のリスト
    • 2次元配列
      • 初期化
        • 始点と終点が同じ場合は0
        • 辺が存在する場合はそのコスト
        • それ以外はINF
  • ループ
    • 3重ループ
    • 1番外側のループがその頂点を経由するかしないか
    • 2番目と3番目のループが始点と終点
    • 経由するかしないかで距離が短いほうで更新

経路復元(ダイクストラ法)

-- 経路復元(ダイクストラ法)
dijkstraPath :: Int -> Int -> Int -> Int -> [(Int, Int, Int)] -> [Int]
dijkstraPath n m start end es = runST $ do
    d <- newArray (0, n - 1) maxBound :: ST s (STArray s Int Int)
    p <- newArray (0, n - 1) (-1) :: ST s (STArray s Int Int)
    writeArray d start 0
    go d p (push (0, start) Empty :: Skew (Int, Int))
    path [end] p
  where
    g = buildGList (0, n - 1) [(a, (b, c)) | (a, b, c) <- es]
    go d p que
        | Just ((c, f), q) <- pop que = do
            rq <- newSTRef q
            df <- readArray d f
            when (c <= df) $ do
                forM_ (g ! f) $ \(t, c) -> do
                    dt <- readArray d t
                    when (df + c < dt) $ do
                        writeArray d t (df + c)
                        writeArray p t f
                        modifySTRef' rq (\q -> push (df + c, t) q)
            go d p =<< readSTRef rq
        | otherwise = return ()
    path (x:xs) p
        | x == start = return (x:xs)
        | otherwise  = do
            pt <- readArray p x
            path (pt:x:xs) p
  • 更新時に手前の頂点をメモっておいて最後に復元する

最小全域木

プリム法

-- プリム法
prim :: Int -> Int -> Int -> [(Int, Int, Int)] -> Int
prim n m start es = runST $ do
    used <- newArray (0, n - 1) False :: ST s (STArray s Int Bool)
    writeArray used start True
    let ss = fmap (\(t, c) -> (c, t)) (g ! start)
    go 0 used (foldr push Empty ss)
  where
    g = buildGList (0, n - 1) [(a, (b, c)) | (a, b, c) <- es]
    go !k used que
        | Just ((c, t), q) <- pop que = do
            usedt <- readArray used t
            if usedt
                then go k used q
                else do
                    let ss = fmap (\(t, c) -> (c, t)) (g ! t)
                    let q' = foldr push q ss
                    writeArray used t True
                    go (k + c) used q'
        | otherwise = return k
  • 用意するもの
    • グラフ(隣接リスト)
    • 優先度付きキュー
      • 初期化:始点から伸びている辺を(コスト, 頂点)の並びで追加
    • 1次元配列
      • 既にその頂点が最小全域木に含まれているかをメモする配列
      • 初期化:始点のみTrue, 残りはFalse
  • ループ
    • キューが空になるまでループ
    • キューから取り出した辺の終点が既に最小全域木に含まれている場合はスキップ
    • そうでなければ採用した辺の終点から伸びている辺を取り出しキューに追加
    • 採用した辺の終点をメモし、再帰

クラスカル法

-- クラスカル法
kruskal :: Int -> Int -> [(Int, Int, Int)] -> Int
kruskal n m es = runST $ do
    uf <- newUF (0, n - 1)
    rc <- newSTRef 0
    forM_ es' $ \(f, t, c) -> do
        isSame <- same uf f t
        when (not isSame) $ do
            unite uf f t
            modifySTRef' rc (+ c)
    readSTRef rc
  where
    es' = sortBy (\(_, _, a) (_, _, b) -> a `compare` b) es
  • 用意するもの
    • 辺のリスト
      • コストの小さい順にソート
    • Union-Find木
  • ループ
    • コストの小さい辺から見ていく
    • 辺の始点と終点が既に同じグループならスキップ
    • そうでなければその辺を採用して始点と終点を同じグループにする

練習問題

初級とか言ってるけど難しすぎでは…?

-- Roadblocks (POJ No.3255)
q2 :: Int -> Int -> [(Int, Int, Int)] -> Int
q2 n m es = runST $ do
    d1 <- newListArray (1, 4) (0 : repeat inf) :: ST s (STUArray s Int Int)
    d2 <- newListArray (1, 4) (0 : repeat inf) :: ST s (STUArray s Int Int)
    go d1 d2 (push (0, 1) Empty)
    readArray d2 n
  where
    inf = 10 ^ 9 + 7
    g = buildGList (1, n) [(f, (t, c)) | (f, t, c) <- es]
    go d1 d2 que
        | Just ((cost, v), que') <- pop que = do
            rque <- newSTRef que'
            d2v <- readArray d2 v
            when (cost <= d2v) $ do
                forM_ (g ! v) $ \(t, c) -> do
                    d1t <- readArray d1 t
                    d2t <- readArray d2 t
                    let d2c = cost + c
                    when (d2c < d1t) $ do
                        writeArray d1 t d2c
                        modifySTRef' rque $ \q -> push (d2c, t) q
                    when (d2c < d2t && d1t < d2c) $ do
                        writeArray d2 t d2c
                        modifySTRef' rque $ \q -> push (d2c, t) q
            go d1 d2 =<< readSTRef rque
        | otherwise = return ()
  • ダイクストラ法で最短距離を求めつつ2番目の最短路も求める
  • キューから取り出して最短距離は更新できないが2番目の最短路が更新できるなら更新
-- Conscription (POJ No.3723)
q3 :: Int -> Int -> Int -> [(Int, Int, Int)] -> Int
q3 n m r es = runST $ do
    uf <- newUF (0, n + m - 1)
    k <- kruskal 0 uf es'
    return $ 10000 * (n + m) + k
  where
    es' = sortBy (\(_, _, c1) (_, _, c2) -> c1 `compare` c2)
        $ concat [[(f, n + t, -c), (n + t, f, -c)] | (f, t, c) <- es]
    kruskal !k uf []             = return k
    kruskal !k uf ((f, t, c):rs) = do
        isSame <- same uf f t
        if isSame
            then kruskal k uf rs
            else do
                unite uf f t
                kruskal (k + c) uf rs
  • 辺のコストを負にして最小全域木を求める
  • なぜこれで答えが求まるのかいまいちわかってない
-- Layout (POJ No.3169)
q4 :: Int -> Int -> Int -> [(Int, Int, Int)] -> [(Int, Int, Int)] -> Int
q4 n ml md as bs = runST $ do
    d <- newListArray (1, n) (0 : repeat maxBound) :: ST s (STUArray s Int Int)
    forM_ [1 .. n] $ \_ -> do
        forM_ as  $ \e -> update d e
        forM_ bs' $ \e -> update d e
        forM_ cs  $ \e -> update d e
    d1 <- readArray d 1
    dn <- readArray d n
    case () of
        _ | d1 < 0         -> return (-1)
          | dn == maxBound -> return (-2)
          | otherwise      -> return dn
  where
    bs' = [(t, f, -c) | (f, t, c) <- bs]
    cs  = [(i + 1, i, 0) | i <- [1 .. n - 1]]
    update d (f, t, c) = do
        df <- readArray d f
        when (df < maxBound) $ do
            dt <- readArray d t
            writeArray d t (min dt (df + c))
  • 普通の最短路問題
    • d(v) + w >= d(u)
      • 頂点vから頂点uへコストwの辺が存在
  • この問題
    • d[i+1] + 0 >= d[i]
      • 頂点i+1から頂点iへコスト0の辺が存在
    • d[AL] + DL >= d[BL]
      • 頂点ALから頂点BLへコストDLの辺が存在
    • d[BD] + (-DD) >= d[AD]
      • 頂点BDから頂点ADへコスト-DDの辺が存在
  • 負の辺が含まれるためベルマン-フォード法で最短経路を求める
  • 負の閉路が存在する場合
    • 例)1と3は仲良しで距離10以内、2と3は仲悪くて距離11以上離す
  • いくらでも離れることが可能な場合
    • 例)Nは誰とも仲良くない

20 Nov 2018, 21:01

Haskellで蟻本(初級編) - データ構造

※蟻本の入力例でしかテストしていません

優先度付きキュー

多くのプログラミング言語では効率的に実装されたプライオリティキューが標準で含まれています。

と書かれているが、Haskellには標準で含まれていない。

じゃあどうするかというと

  1. containersMapSetMin/Maxの関数で頑張る
  2. 自分で作る

の2択になると思う。


Leftist Heap

module Heap
    ( Heap (..)
    , meld
    , push
    , pop
    ) where

import Control.Monad

data Heap a
    = Empty
    | Node !Int !a !(Heap a) !(Heap a)
    deriving Show

rank :: Heap a -> Int
rank Empty          = 0
rank (Node s _ _ _) = s

heap :: a -> Heap a -> Heap a -> Heap a
heap x a b
    | ra < rb   = Node (ra + 1) x b a
    | otherwise = Node (rb + 1) x a b
  where
    ra = rank a
    rb = rank b

meld :: Ord a => Heap a -> Heap a -> Heap a
meld a     Empty = a
meld Empty b     = b
meld a@(Node _ ax al ar) b@(Node _ bx bl br)
    | ax < bx   = heap ax al (meld ar b)
    | otherwise = heap bx bl (meld br a)

push :: Ord a => a -> Heap a -> Heap a
push x a = meld a (Node 1 x Empty Empty)

pop :: Ord a => Heap a -> Maybe (a, Heap a)
pop Empty          = Nothing
pop (Node _ x l r) = Just (x, meld l r)

test :: (Show a, Ord a) => [a] -> IO ()
test xs = do
    heap <- flip (flip foldM Empty) xs $ \h i -> do
        let h' = push i h
        print h'
        return h'
    putStrLn "--------------------"
    flip (flip foldM_ heap) [1 .. length xs] $ \h _ -> do
        print h
        let Just (_, h') = pop h
        return h'
*Heap> test [4,1,2,5,3]
Node 1 4 Empty Empty
Node 1 1 (Node 1 4 Empty Empty) Empty
Node 2 1 (Node 1 4 Empty Empty) (Node 1 2 Empty Empty)
Node 2 1 (Node 1 4 Empty Empty) (Node 1 2 (Node 1 5 Empty Empty) Empty)
Node 2 1 (Node 2 2 (Node 1 5 Empty Empty) (Node 1 3 Empty Empty)) (Node 1 4 Empty Empty)
--------------------
Node 2 1 (Node 2 2 (Node 1 5 Empty Empty) (Node 1 3 Empty Empty)) (Node 1 4 Empty Empty)
Node 2 2 (Node 1 5 Empty Empty) (Node 1 3 (Node 1 4 Empty Empty) Empty)
Node 2 3 (Node 1 4 Empty Empty) (Node 1 5 Empty Empty)
Node 1 4 (Node 1 5 Empty Empty) Empty
Node 1 5 Empty Empty

Skew Heap

  • Leftist Heapから木のバランスを取る処理を除いたもの
  • 実装が軽いのに速さもそんな変わらないというのをどこかで見た
  • 大きい順に取り出したいときは要修正
module Skew
    ( Skew (..)
    , meld
    , push
    , pop
    ) where

import Control.Monad

data Skew a
    = Empty
    | Node !a !(Skew a) !(Skew a)
    deriving Show

meld :: Ord a => Skew a -> Skew a -> Skew a
meld a     Empty = a
meld Empty b     = b
meld a@(Node ax al ar) b@(Node bx bl br)
    | ax < bx   = Node ax al (meld ar b)
    | otherwise = Node bx bl (meld br a)

push :: Ord a => a -> Skew a -> Skew a
push x a = meld a (Node x Empty Empty)

pop :: Ord a => Skew a -> Maybe (a, Skew a)
pop Empty        = Nothing
pop (Node x l r) = Just (x, meld l r)

test :: (Show a, Ord a) => [a] -> IO ()
test xs = do
    skew <- flip (flip foldM Empty) xs $ \h i -> do
        let h' = push i h
        print h'
        return h'
    putStrLn "--------------------"
    flip (flip foldM_ skew) [1 .. length xs] $ \h _ -> do
        print h
        let Just (_, h') = pop h
        return h'
*Skew> test [4,1,2,5,3]
Node 4 Empty Empty
Node 1 Empty (Node 4 Empty Empty)
Node 1 Empty (Node 2 Empty (Node 4 Empty Empty))
Node 1 Empty (Node 2 Empty (Node 4 Empty (Node 5 Empty Empty)))
Node 1 Empty (Node 2 Empty (Node 3 Empty (Node 4 Empty (Node 5 Empty Empty))))
--------------------
Node 1 Empty (Node 2 Empty (Node 3 Empty (Node 4 Empty (Node 5 Empty Empty))))
Node 2 Empty (Node 3 Empty (Node 4 Empty (Node 5 Empty Empty)))
Node 3 Empty (Node 4 Empty (Node 5 Empty Empty))
Node 4 Empty (Node 5 Empty Empty)
Node 5 Empty Empty

プライオリティキューを用いる問題

-- Expedition (POJ 2431)
q1 :: Int -> Int -> Int -> [Int] -> [Int] -> Int
q1 n l p as bs = go 0 p Empty (zip bs as)
  where
    go !k fuel pque bas
        | fuel >= l    = k
        | Just ((g, _), pque'') <- pop pque'
            = go (k + 1) (fuel + g) pque'' ys
        | otherwise = -1
      where
        (xs, ys) = span ((<= fuel) . snd) bas
        pque' = foldr push pque xs
  • fuelがゴールまで足りていたら補給回数を返す
  • そうでなければ通り過ぎたガソリンスタンドを全てキューに追加し、補給できる量が最大のものを取り出す
  • 取り出せた場合は再帰、取り出せなかったら-1
-- Fence Repair (POJ 3253)
q2 :: Int -> [Int] -> Int
q2 n ls = go 0 (foldr push Empty ls)
  where
    go !k s
        | Just (x1, s1) <- pop s
        , Just (x2, s2) <- pop s1
            = let x = x1 + x2 in go (k + x) (push x s2)
        | otherwise = k

Union-Find木

  • 経路圧縮のみ
module UF
    ( UF (..)
    , newUF
    , root
    , same
    , unite
    ) where

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

newtype UF s
    = UF (STUArray s Int Int)

newUF :: (Int, Int) -> ST s (UF s)
newUF (i, j) = UF <$> (newListArray (i, j) [i .. j])

root :: UF s -> Int -> ST s Int
root (UF a) i = do
    p <- readArray a i
    if p == i
        then return p
        else do
            rp <- root (UF a) p
            writeArray a i rp
            return rp

same :: UF s -> Int -> Int -> ST s Bool
same uf x y = (==) <$> (root uf x) <*> (root uf y)

unite :: UF s -> Int -> Int -> ST s ()
unite uf@(UF a) i j = do
    ri <- root uf i
    rj <- root uf j
    when (ri /= rj) $ do
        writeArray a rj ri

Union-Find木を用いる問題

-- 食物連鎖 (POJ 1182)
q3 :: Int -> Int -> [(Int, Int, Int)] -> Int
q3 n k qs = runST $ do
    uf <- newUF (1, n * 3)
    go 0 uf qs
  where
    go !f _  [] = return f
    go !f uf ((1, x, y):rs) = do
        let p1 = or [x < 1, y < 1, x > n, y > n]
        p2 <- (||) <$> same uf x (y + n) <*> same uf x (y + n * 2)
        if p1 || p2
            then go (f + 1) uf rs
            else do
                unite uf x y
                unite uf (x + n) (y + n)
                unite uf (x + n * 2) (y + n * 2)
                go f uf rs
    go !f uf ((2, x, y):rs) = do
        let p1 = or [x < 1, y < 1, x > n, y > n]
        p2 <- (||) <$> same uf x y <*> same uf x (y + n * 2)
        if p1 || p2
            then go (f + 1) uf rs
            else do
                unite uf x (y + n)
                unite uf (x + n) (y + n * 2)
                unite uf (x + n * 2) y
                go f uf rs
  • n = 100のときUnion-Findの要素数はその3倍として

    • 1: 1が種類A, 101: 1が種類B, 201: 1が種類C
    • 2: 2が種類A, 102: 2が種類B, 202: 2が種類C
    • ...
  • 与えられた情報より、同時に起こるものを同じグループとしていく

    • 情報1: 1と2は同じ種類です。
      • unite uf 1 2, unite uf 101 102, unite uf 201 202
    • 情報2: 1と2を食べます。
      • unite uf 1 102, unite uf 101 202, unite uf 201 2

19 Nov 2018, 22:09

Haskellで蟻本(初級編) - 動的計画法

※蟻本の入力例でしかテストしていません

{-# LANGUAGE BangPatterns, FlexibleContexts #-}

import Control.Monad
import Data.List
import qualified Data.ByteString.Char8 as B

import Control.Monad.ST
import Data.Array.ST
import Data.Array
import qualified Data.Map.Strict as M

-- 01 ナップサック問題
q1 :: Int -> Int -> [(Int, Int)] -> Int
q1 n m wvs = go n m
  where
    wva = listArray (1, n) wvs
    memo = listArray ((0, 0), (n, m)) [go i j | i <- [0 .. n], j <- [0 .. m]]
    go 0 _ = 0
    go i j
        | j < w     = x1
        | otherwise = max x1 x2
      where
        (w, v) = wva ! i
        x1 = memo ! (i - 1, j)
        x2 = memo ! (i - 1, j - w) + v

q1' :: Int -> Int -> [(Int, Int)] -> Int
q1' n m wvs = runST $ do
    dp <- newArray ((0, 0), (n, m)) 0 :: ST s (STUArray s (Int, Int) Int)
    forM_ [1 .. n] $ \i -> do
        let (w, v) = wva ! i
        forM_ [0 .. m] $ \j -> do
            x1 <- readArray dp (i - 1, j)
            if j < w
                then writeArray dp (i, j) x1
                else do
                    x2 <- readArray dp (i - 1, j - w)
                    writeArray dp (i, j) (max x1 (x2 + v))
    readArray dp (n, m)
  where
    wva = listArray (1, n) wvs

q1'' :: Int -> Int -> [(Int, Int)] -> Int
q1'' n m wvs = runST $ do
    dp   <- newArray ((0, 0), (n, m)) 0     :: ST s (STUArray s (Int, Int) Int)
    memo <- newArray ((0, 0), (n, m)) False :: ST s (STUArray s (Int, Int) Bool)
    sequence_ [writeArray memo (0, i) True | i <- [0 .. m]]
    go dp memo (n, m)
  where
    wva = listArray (1, n) wvs
    go dp memo (i, j) = do
        exists <- readArray memo (i, j)
        when (not exists) $ do
            let (w, v) = wva ! i
            x1 <- go dp memo (i - 1, j)
            if j < w
                then do
                    writeArray dp (i, j) x1
                else do
                    x2 <- go dp memo (i - 1, j - w)
                    writeArray dp (i, j) (max x1 (x2 + v))
            writeArray memo (i, j) True
        readArray dp (i, j)
  • メモ化再帰 これでAC取りたい
  • 動的計画法 圧倒的速さ
  • メモ化再帰 上のでTLEになるがメモ化再帰したいとき?
-- 最長共通部分列問題
q2 :: Int -> Int -> B.ByteString -> B.ByteString -> Int
q2 n m ss ts = go (n, m)
  where
    memo = listArray ((0, 0), (n, m)) [go (i, j) | i <- [0 .. n], j <- [0 .. m]]
    go (0, _) = 0
    go (_, 0) = 0
    go (i, j)
        | si == tj  = maximum [x1, x2, x3]
        | otherwise = max x1 x2
      where
        si = ss `B.index` (i - 1)
        tj = ts `B.index` (j - 1)
        x1 = memo ! (i - 1, j)
        x2 = memo ! (i, j - 1)
        x3 = memo ! (i - 1, j - 1) + 1
  • ByteStringindexO(1)
  • where句に細かく名前つけてあげると読みやすくなる気がする
-- 個数制限なしナップサック問題
q3 :: Int -> Int -> [(Int, Int)] -> Int
q3 n m wvs = go (n, m)
  where
    wva  = listArray (1, n) wvs
    memo = listArray ((0, 0), (n, m)) [go (i, j) | i <- [0 .. n], j <- [0 .. m]]
    go (0, _) = 0
    go (i, j)
        | j < w     = x1
        | otherwise = max x1 x2
      where
        (w, v) = wva ! i
        x1 = memo ! (i - 1, j)
        x2 = memo ! (i, j - w) + v
  • このように計算量を減らすのだって書いてあるけど思いつけるようになるのだろうか
-- 01 ナップサック問題 その2
q4 :: Int -> Int -> [(Int, Int)] -> Int
q4 n m wvs = maximum $ do
    v <- [0 .. vmax]
    let w = memo ! (n, v)
    guard (w <= m)
    return v
  where
    inf = 10 ^ 10
    vmax = sum $ fmap snd wvs
    wva = listArray (1, n) wvs
    memo = listArray ((0, 0), (n, vmax)) [go (i, j) | i <- [0 .. n], j <- [0 .. vmax]]
    go (0, 0) = 0
    go (0, _) = inf
    go (i, j)
        | j < v     = x1
        | otherwise = min x1 x2
      where
        (w, v) = wva ! i
        x1 = go (i - 1, j)
        x2 = go (i - 1, j - v) + w
-- 個数制限つき部分和問題
q5 :: Int -> Int -> [Int] -> [Int] -> Bool
q5 n k as ms
    | go (n, k) < 0 = False
    | otherwise     = True
  where
    aa = listArray (1, n) as
    ma = listArray (1, n) ms
    memo = listArray ((0, 0), (n, k)) [go (i, j) | i <- [0 .. n], j <- [0 .. k]]
    go (0, 0) = 0
    go (0, _) = -1
    go (i, j)
        | x1 >= 0          = m
        | j < a || x2 <= 0 = -1
        | otherwise        = x2 - 1
      where
        a = aa ! i
        m = ma ! i
        x1 = memo ! (i - 1, j)
        x2 = memo ! (i, j - a)
-- 最長増加部分列問題
q6 :: Int -> [Int] -> Int
q6 n as = go n
  where
    aa = listArray (1, n) as
    memo = listArray (1, n) [go i | i <- [1 .. n]]
    go 1 = 1
    go i = maximum . (1:) $ do
        j <- [1 .. i - 1]
        let (ai, aj) = (aa ! i, aa ! j)
        guard (aj < ai)
        return $ (memo ! j) + 1

q6' :: Int -> [Int] -> Int
q6' n as = maximum $ foldl' f M.empty as
  where
    f m a
        | Just (k, v) <- M.lookupGE a m = M.insert a v $ M.delete k m
        | otherwise                     = M.insert a (M.size m + 1) m
  • IntMapと異なりMapsizeO(1)
  • 蟻本の「DPテーブルの変化」の図と同じようにMapが更新されていく
-- 分割数
q7 :: Int -> Int -> Int -> Int
q7 n m d = go (m, n)
  where
    memo = listArray ((0, 0), (m, n)) [go (i, j) | i <- [0 .. m], j <- [0 .. n]]
    go (0, 0) = 1
    go (0, _) = 0
    go (i, j)
        | j < i     = x1
        | otherwise = (x1 + x2) `mod` d
      where
        x1 = memo ! (i - 1, j)
        x2 = memo ! (i, j - 1)
漸化式がよく理解できなかったのでメモ

dp[3][4] = 4 (4の3分割の総数)

  (4, 0, 0)
  (3, 1, 0) -- dp[2][4]と同じ
  (2, 2, 0)

  (2, 1, 1) -- 4の3分割 (全て1以上) == (4 - 3)の3分割 (全て0以上)
-- 重複組合せ
q8 :: Int -> Int -> Int -> [Int] -> Int
q8 n m d as = go (n, m)
  where
    aa = listArray (1, n) as
    memo = listArray ((0, 0), (n, m)) [go (i, j) | i <- [0 .. n], j <- [0 .. m]]
    go (0, 0) = 1
    go (0, _) = 0
    go (i, j)
        | j - 1 < 0      = x1
        | j - 1 - ai < 0 = (x1 + x2) `mod` d
        | otherwise      = (x1 + x2 - x3) `mod` d
      where
        ai = aa ! i
        x1 = memo ! (i - 1, j)
        x2 = memo ! (i, j - 1)
        x3 = memo ! (i - 1, j - 1 - ai)