ABC411をHaskellで

A - Required Length

問題 ABC411A

シグネチャを決める。

abc411a :: String -- P
        -> Int    -- L
        -> Bool   -- 答え
abc411a p l = l <= length p

-- 別解
abc411a p l = not $ null $ drop (pred l) p
-- より素直な別解
abc411a p l = length (take l p) == l

Pがとても長い場合もあるとき、前者だと $O(|P|)$ かかるが、後者だと $O(\min(L,|P|))$ で済む。

B - Distance Table

問題 ABC411B

シグネチャを決める。
答えは、順にそれぞれの行に出力するべき値のリストとする。

abc411b :: Int     -- N
        -> [Int]   -- Di
        -> [[Int]] -- 答え

$i$ 行めには、駅 $i$ を起点としての距離の累積和を与える。

結果

import Data.List

abc411b :: Int -> [Int] -> [[Int]]
abc411b n ds = take (pred n) $ map (scanl1 (+)) $ tails ds

C - Black Intervals

問題 ABC411C

シグネチャを決める。

abc411c :: Int   -- N
        -> Int   -- Q
        -> [Int] -- Ai
        -> [Int] -- 答え

マス0とマス$N+1$も存在して、常に白いとみなす。
左右を含めた状況を列挙して、区間の個数の変化を考える。

マス $A-1$	マス $A$	マス $A+1$	増減
黒	白→黒	黒	$-1$
白	白→黒	黒	$\pm 0$
黒	白→黒	白	$\pm 0$
白	白→黒	白	$+ 1$
黒	黒→白	黒	$+1$
白	黒→白	黒	$\pm 0$
黒	黒→白	白	$\pm 0$
白	黒→白	白	$- 1$

つまり、左右が同じ色なとき、同じ色にするなら $-1$ 違う色にするなら $+1$ それ以外のときは変化なし。

色の現状はmutable arrayで管理する。

結果

黒を True として Vector.Mutable でマスの状態を追跡する。

import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad.ST

abc411c :: Int -> Int -> [Int] -> [Int]
abc411c n _q as = runST $
  do
    fld <- MUV.replicate (n + 2) False :: ST s (MUV.MVector s Bool)
    scanl1 (+) <$> forM as (\a -> do
      xyz <- mapM (MUV.read fld) [pred a .. succ a]
      let x:y:z:_ = xyz
      MUV.write fld a (not y)
      return $ if x /= z then 0 else if x == y then 1 else -1
      )

もっと $A_i$ の範囲が広くて配列では持てなくなったら、MEX用のデータ構造を使うことになるだろうか。

D - Conflict 2

問題 ABC411D

シグネチャを決める。
タイプ2以外のクエリでは文字列は無視する。

type Query = (Char,Int,String)

abc411d :: Int      -- N
        -> Int      -- Q
        -> [Query]  -- query_i
        -> String   -- 答え

サーバを0番として、$N+1$ 要素の配列に状況を保持する。
ただし、末尾に連結すると遅いので文字列のリストの形で持っておき、最後に連結する。

結果

import qualified Data.Vector.Mutable as MV
import Control.Monad.ST
import Control.Monad

abc411d :: Int -> Int -> [Query] -> String
abc411d n _q qs = runST $ do
  v <- MV.replicate (succ n) [] :: ST s (MV.MVector s [String])
  forM_ qs $ \(t,p,s) -> do
    case t of
      '1' -> MV.read v 0 >>= MV.write v p
      '2' -> MV.modify v s p
      '3' -> MV.read v p >>= MV.write v 0
  v0 <- MV.read v 0
  return $ concat $ reverse v0

クエリを Query 型に読み込むところの方が面倒な仕事かもしれない。

E - E [max]

問題 ABC411E

シグネチャを決める。

abc411e :: Int      -- N
        -> [[Int]]  -- Aij
        -> Int      -- 答え

自分の考え

期待値はスコアとその確率の積の総和なので、全てのサイコロの目の値について、その目が最大値となる確率を求めることが目標となる。

ここでのサイコロの目を全て集めて昇順に背番号を付けて $X_1, \dots, X_M$ と呼ぶ。また $X_0 = 0$ とする。

最大値が $X_i$ になる確率は
（全てのサイコロの出目が $X_i$ 以下になる確率）－（全てのサイコロの出目が $X_i$ 未満になる確率）
ここで、サイコロの出目は離散的なので、後者は（全てのサイコロの出目が $X_{i-1}$ 以下になる確率）と読み替えられる。

「全てのサイコロの出目が $X_i$ 以下になる確率」を $P(X_i)$ とする。
問題の要求である期待値は $\sum_i X_i \big \{P(X_i) - P(X_{i-1}) \big \}$ である。

それぞれのサイコロ $j$ について、$X_i$ 以下の目の面数を $M_j[X_i]$ とすると
$P(X_i) = \prod_j (M_j[X_i]/6)$
となる。しかしこれを直接数えるのは難しい。
$M_j[X_i]$ は、サイコロ $j$ が目 $X_i$ を $k$ 面持つときにそれだけ増加する。$M_j[X_i] = M_j[X_{i-1}] + k$

サイコロは6面なので $0 \leq M_j[X_i] \leq 6$ である。
$N$個のサイコロ全体で、$M_j[X_i] = k$ となるサイコロ $j$ の個数を $K(X_i)[k]$ とする。または長さ7のベクトルで書く。
これは全て足すと$N$になる。$\sum_k K(X_i)[k] = N$
0以下の面を持つサイコロは存在しないので $K(X_0) = [N,0,0,0,0,0,0]$
あるサイコロが目 $X_i$ 未満の面を $a$ 面、目 $X_i$ の面を $b$ 面もつとき、$K(X_i)$ は第 $a$ 要素を1減らし、第 $a+b$ 要素を1増やすことになる。

$K$を用いると $\displaystyle P(X_i) = \prod_{0 \leq k \leq 6} \big ( \frac{k}{6} \big )^{K(X_i)[k]}$ と書ける。

結果

import qualified Data.IntMap as IM
import qualified Data.Vector.Unboxed as UV

abc411e :: Int -> [[Int]] -> Int
abc411e n ass = summ ans
  where
    vd1 = UV.fromListN 7 [-1,1,0,0,0,0,0] -- K[k]をひとつずらす差分
    vd2 = UV.fromListN 7 [0,-1,1,0,0,0,0]
    vd3 = UV.fromListN 7 [0,0,-1,1,0,0,0]
    vd4 = UV.fromListN 7 [0,0,0,-1,1,0,0]
    vd5 = UV.fromListN 7 [0,0,0,0,-1,1,0]
    vd6 = UV.fromListN 7 [0,0,0,0,0,-1,1]
    im :: IM.IntMap (UV.Vector Int)       -- XiごとにK[]の差分を寄せて累積和をとる
    im = IM.fromListWith (UV.zipWith (+)) $ concat
         [ zip as1 [vd1,vd2,vd3,vd4,vd5,vd6]
         | as <- ass, let as1 = sort as]

    r6 = modRecip 6                     -- 1/6
    ps = take 7 $ iterate (add r6) 0    -- 0/6, 1/6, ..., 6/6
    v0 = UV.fromListN 7 [n,0,0,0,0,0,0] -- K(0)
    (_vN, ans) = mapAccumL step (v0, 0) $ IM.assocs im
    step :: (UV.Vector Int              -- K(X_{i-1})[]
            ,Int)                       -- P(X[i-1])
         -> (Int, UV.Vector Int)        -- XiとKの差分
         -> ((UV.Vector Int,Int), Int)  -- K(X_i), P(Xi), 期待値の成分
    step (kv, p0) (x, dv) = ((kv1, p1), mul x $ reg $ p1 - p0)
      where
        kv1 = UV.zipWith (+) kv dv
        p1 = prodd $ zipWith mpower ps (UV.toList kv1)

-- モジュロ演算は省略

問題の内容に対してほぼストレートに計算するアプローチで、制限時間 3sec を 1946ms, 110MB でACした。

アライさんのヒント

フレンズさんいわく

アライグマ「E問題は主客転倒なのだ！　期待値は「1以上になる確率＋2以上になる確率＋3以上になる確率＋…」だから、サイコロの目それぞれについて「X以上になる確率」を求めてうまく足せばいいのだ！」

この逆転は ABC295E の解説
https://atcoder.jp/contests/abc295/editorial/6048
に書かれているけど、公式解説では採用されていない感じ？

すごい方法

他の人の結果を眺めていて、めちゃ短くてめちゃ速い提出727byte, 313ms, 11MBを見つけた。何だこれ？

読み解くと、

• サイコロの背番号をキーにしたセグメント木に、そのサイコロが注目している値未満の目を出す確率、つまり $k/6$ を保持
• セグメント木の演算は掛け算、単位元は1
• （サイコロの目 a，背番号 i）という対のリストを小さい順に調べて、
  • サイコロ i がちょうどこの目 a （重複があったとしても、今注目しているこの a の面）を出して、それが最大値となる確率 に a を掛けた期待値の成分を答えに足し込む
  • 確率は $(\text{1番からi-1番までのサイコロの確率の積}) \cdot 1/6 \cdot (\text{i+1番からN番までのサイコロの確率の積})$
  • セグメント木の i 番の確率を 1/6 だけ増加させる

という方針だった。これはすごいので写経する。

import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Algorithms.Intro as VAI

abc411e :: Int -> [[Int]] -> Int
abc411e n ass = runST $
  do
    st <- makeRQArray mul 1 $ replicate n 0
    UV.foldM (\acc (a, i) -> do
      vi <- readRQArray st i
      qi <- writeRQArray st i recip6
      writeRQArray st i $ add vi recip6
      return $! add acc $ mul a qi
      ) 0 aiV
  where
    aiV :: UV.Vector (Int, Int)
    aiV = UV.modify (VAI.sortBy compare) $ UV.fromListN (n * 6) [(a, i) | (i, as) <- zip [0 ..] ass, a <- as]
    recip6 = modRecip 6

-- セグメント木は省略

オリジナルは、セグメント木に区間の問いあわせ ($O(\log N)$) を2度、背番号 i の値を問いあわせ ($O(1)$)、$+1/6$ した値を書き戻して木を更新 $O(\log N)$ としているところを、
背番号 i の値を問いあわせ ($O(1)$)、背番号 i の値を $1/6$ に強制上書きして木を更新しつつ、木全体の結果を問いあわせ ($O(\log N)$)、$+1/6$ した値を書き戻して木を更新 $O(\log N)$ として、係数を下げる工夫を加えた。
木を更新する手順の完了時には根の値が副産物として手に入り、それがちょうど使いたい値なので利用した。

提出結果：343ms, 61MB

F - Contraction

問題 ABC411F

シグネチャを決める。横着する。

abc411f :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- Ui, Vi
        -> Int      -- Q
        -> [Int]    -- Xi
        -> [Int]    -- 答え

自分の考え

コマとは、初期状態で頂点に接続していた辺が、縮約操作で頂点の番号が変えられても、どこに接続しているか追跡できるようにする目印。むしろヒント。

縮約により2頂点が一つに同一視されていくので、この過程を追跡して状況を取り出せるようにするには、Union-Findが使えそう。
元の頂点番号、もしはコマの番号をキーに、Union-Findの代表元を、縮約した結果の現在の現在の頂点番号として使う。

このとき、実行するべき操作は順に

1. クエリで指定された辺の番号 $X_i$ から、元の両端の頂点番号 $U_i, V_i$ を取り出す（配列で $O(1)$）
2. 現在のグラフにおけるその番号のコマの位置 $A_i, B_i$ を求める（Union-Findの代表元を読み出す）
3. $A_i = B_i$ なら、この縮約は無効なので辺の本数は変化しない。終わり。
4. $A_i$ と $B_i$ を Union-Find で一つにする。新たな代表元を $C_i$ とする
5. $A_i$ と $B_i$ を結んでいた辺を除く
6. それ以外の頂点と $A_i, B_i$ を結んでいた辺を全て $C_i$ に結び直す
7. このとき $A_i, B_i$ の両方ともに辺を持っていた頂点の個数、もしくは $A_i$ の隣接頂点と$B_i$ の隣接頂点の共通部分のサイズだけ、辺が重なることで減る
8. 辺の本数を更新して終了

手順4において、$C_i$ を全く別の番号にすると、以降の操作が常に隣接頂点全てにかかるので重い。
いずれかの頂点にするべきである。それはUnion-Findの動作とも一致する。
そしていずれかにすることで、以降の操作が全て、代表でなくなる方の頂点に冠するものだけにできる。
これをなるべく減らす方がよい。具体的には、隣接頂点が少ない方にするべきである。

これをすると、大きい方はあまり動かず、小さい方が大きい方に飲み込まれる（と、大きい方はますます変化しなくなる）ことで、全体でみると計算量がかなり節約できる。
どうやらこの発想を「マージテク」と界隈では呼んでいるようだ。

アライグマ「F問題はマージテクなのだ！　頂点をマージするとき、「コマの個数＋辺の個数」が多い方にマージすることにすればいいのだ！」

自分のUnion-Findは、代表元を求める getRoot 操作により経路圧縮を行うため、代表元を選び直す操作において木のランクなどは気にせず、引数の順で固定している。この仕様がこの問題ではうまく使えた。

結果

頂点に対して、隣接頂点集合を更新できる配列で持つ必要がある。
これは、そのサイズを高速に取り出せる必要があるので、IntSet でなく Set Int を選ぶ。

import Control.Monad
import Data.List

import Control.Monad.ST
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed.Mutable as MUV
import qualified Data.Set as S
import qualified Data.Vector.Unboxed as UV

abc411f :: Int -> Int -> [[Int]] -> Int -> [Int] -> [Int]
abc411f n m uvs q xs = runST $
  do
-- xi に対して U,V を取り出す配列
    let e = UV.fromListN (succ m) $ (0,0) : [(u,v) | u:v:_ <- uvs]
-- 接続辺番号をSetで持つグラフ配列（可変）
    g <- MV.replicate (succ n) S.empty
    UV.forM_ (UV.tail e) (\(u,v) -> do
      MV.modify g (S.insert v) u
      MV.modify g (S.insert u) v
      )
-- Union-Find
    uf <- newUF (succ n)
-- 答えを入れる配列
    ans <- MUV.new (succ q) :: ST s (MUV.MVector s Int)
    MUV.write ans 0 m
-- クエリに順に対応
    forM_ (zip [1 ..] xs) (\(i, x) -> do
      ans0 <- MUV.read ans (pred i)
      let (u, v) = e UV.! x
      a <- getRoot uf u
      b <- getRoot uf v
      if a == b then MUV.write ans i ans0 else do
        sa <- S.delete b <$> MV.read g a
        sb <- S.delete a <$> MV.read g b
        let (c,d,sc,sd) = if S.size sa <= S.size sb then (a,b,sa,sb) else (b,a,sb,sa)
        MUV.write uf c d -- uniteUF uf c d
        MV.write g c S.empty
        MV.write g d $ S.union sc sd
        ans1 <- foldM (\ansi y -> do
          gy <- MV.read g y
          MV.write g y $ S.insert d $ S.delete c gy
          return $! if S.member d gy then pred ansi else ansi
          ) (pred ans0) (S.elems sc)
        MUV.write ans i ans1
      )
    forM [1 .. q] (MUV.read ans)

-- Union-Findは省略

公式解説の方法

上ではコマの現在位置を Union-Find で管理したが、これをより原始的な、コマ番号から現在の頂点番号を取り出せるだけの更新できる配列で管理する。
すると、代表元とならなかった方の頂点に乗っているコマ全員について、この表の更新の手間がかかる。（上の手順4がUnion-Findのunion操作から配列の更新に置き換わる。）

代表元を選ぶ根拠として、隣接頂点の個数だけでなく、コマの人数も勘案することで、この場合に有利な向きを選択する。

これの方が計算量は大きいが、プログラムは単純になる。少なくとも命令型では。
書き換え可能配列をゴリゴリするため、Haskellで書いてもしっくりこなかった。

G - Count Cycles

問題 ABC411G

シグネチャを決める。横着する。

abc411g :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- Ui, Vi
        -> Int      -- 答え

$N \leq 20$ とごく小規模なグラフに見えるが、$M \leq 2 \times 10^5$ で多重辺だらけなので、ちゃんと高速な読み込みを使おう。

全然わからないのでフレンズさんのヒントを見てもまだわからなくて、公式解説を見た。
一度のDPで全てを計算できる訳ではなくて、最大の頂点番号sを切り替えて何度もDPを計算した総和を求める必要があるというのがミソのようだ。

遅延配列による集めるDP、純粋計算

Haskellらしい（と僕が勝手に思っている）やりかた。
変数を定義する。

• 頂点 $i,j$ 間の辺の本数 $C_{i,j}$
• 今考えている最大の頂点番号 $s$
• DP配列 $cnt[S][j]$ 頂点 $s$ から始まり、頂点集合 $S$ に含まれる頂点をちょうど一度ずつ通って、頂点 $j$ に至るパスの本数

すると、次のようになる。なお、常に $s \in S$ である。

cnt[S][j] =
\begin{cases}
0 & (j \not\in S) \\
1 & (j = s, S = \{s\}) \\
0 & (j = s) \\
\sum_{i \in S} C_{i,j} \times cnt[S \setminus \{j\}][i] & (j \neq s)
\end{cases}

$s$ から $i$ まで進み、さらに $j$ に進むやり方を全て数えている。

頂点集合 $S$ の $s$ から始めて $j$ で終わる経路に、さらに $s$ へ戻る方法が $C_{j,s}$ とおりあるので、これを合計すれば答えとなる。ただし反対回りのものも数えて倍になっているので半分にする。あと、長さ2のサイクルは別に数える。

$\sum_{i<j} C_{i,j} \ (C_{i,j} - 1)/2 + \frac{1}{2} \sum_s \sum_{S,j} cnt[S][j] \times C_{j,s}$

結果(MLE)

import Data.Array
import Data.Bits
import Data.List
import Control.DeepSeq

abc411g :: Int -> Int -> [[Int]] -> Int
abc411g n _m uvs = mul (modRecip 2) (summ $ twos2 : map compute [2 .. n1])
  where
    n1 = pred n
    g = accumArray (+) 0 ((0,0), (n1,n1)) $ -- 説明の C_ij
        [((pred u, pred v),1) | u:v:_ <- uvs] ++
        [((pred v, pred u),1) | u:v:_ <- uvs]
-- 長さ2のサイクル
    twos2 = sum [k * pred k | ((i,j), k) <- assocs g, i < j, k > 1]
-- 長さ3以上のサイクルについてDPする
    compute s = deepseq cnts $ summ [mul c $ g ! (i, s) | ((bs, i), c) <- assocs cnts, popCount bs >= 3 ]
      where
        bnds = ((bit s, 0), (pred $ bit $ succ s, s)) -- s は必ず bs の最上位ビット
        cnts :: Array (Int,Int) Int
        cnts = listArray bnds $ map f $ range bnds
        f (bs, j)
          | not $ testBit bs j = 0 -- j \not\in S は無効
          | bs == bit s = 1 -- [{s}][s] たね
          | j == s = 0 -- 先頭と末尾がsで一致するのは {s} 以外無理
          | otherwise = summ [mul (g ! (i, j)) (cnts ! (bsj, i)) | let bsj = clearBit bs j, i <- [0 .. s], testBit bsj i]

-- モジュロ演算は省略

提出：ACx26, MLEx14, 5598ms, 1502MiB
6secという時間制限はクリアできたが、1024MiB というメモリ制限から足が出てしまう。

配るDPをなるべく純粋に

集合のビット表現を配列の添え字に使うDPは、部分集合の方が小さな値になるので、添え字の順に計算すると、部分集合に対する結果は計算済みになっている、という性質の使い方をする。

そして配るDPは、それより手前の計算を通して足し込まれた値が結果として確定され、より添え字の大きい範囲に値を配る動作を繰り返す。このうち、足し込むという動作は書き換え的だが、配る値と配る先を決める計算は純粋にも書けるので、そこを括り出す。

計算結果はimmutableな配列に戻すのが望ましいが、UArray にしようとするとうまくいかなかったので、Vector.Unboxed に Ix で添え字を付けた擬似配列にいれて返すことにした。

import Control.Monad.ST
import qualified Data.Vector.Unboxed.Mutable as MUV
import qualified Data.Vector.Unboxed as UV
import Data.Ix

-- 擬似配列
-- readPA (PArray i a) i で a 型の要素が読み出される
-- keys, elems, assocs などの関数も同様にして生やせる
-- が、素直にbndsとvecの対にしておくべきかも。

newtype PArray i a = PA {readPA :: i -> a}

makePArray :: (UV.Unbox a, Ix i) => (i, i) -> UV.Vector a -> PArray i a
makePArray bnds vec = PA (\i -> vec UV.! index bnds i)

-- 配るDPを実行する
distribute :: (Ix i, MUV.Unbox a)
           => (i, i)              -- 添え字の範囲
           -> a                   -- 配列のデフォルト値
           -> (b -> a -> a)       -- 配られた値を足し込む演算
           -> [(i, a)]            -- 計算の種になる初期値
           -> (i -> a -> [(i,b)]) -- 添え字と確定した値を使って、さらに配る先と配る値のリストを作る関数 配る先は現在位置より後方限定
           -> PArray i a          -- 結果 擬似配列
distribute bnds zero add inis f = runST $ do
  vec <- MUV.replicate (rangeSize bnds) zero
  mapM_ (\(i, a) -> MUV.write vec (index bnds i) a) inis
  mapM_ (\i -> do
    v <- MUV.read vec (index bnds i)
    mapM_ (\(j,u) ->
      MUV.modify vec (add u) (index bnds j)
      ) (f i v)
    ) (range bnds)
  makePArray bnds <$> UV.freeze vec

accumArray の変種という感じだろうか。

これを用いて、上の compute だけ、配るDPで作り直す。
公式解説に
初期化：$dp[\{s\}][s] = 1$ とある初期値は inis に [((bit s, s), 1)] と指定し、
遷移：$dp[S\cup\{j\}][j] \leftarrow \text{(左辺)} + dp[S][i] \times C_{i,j} (j \not\in S)$ は、足し込みを (+) とし、$[S][i]$ の値が決まったとき、各 $j$ に関して、$\times C_{i,j}$ した値を $[S\cup\{j\}][j]$ に配る動作を f で指定する。

    compute s = summ [mul v $ g ! (i, s) | bsi@(bs,i) <- range bnds, popCount bs >= 3, let v = readPA arr bsi]
      where
        bnds = ((bit s, 0), (pred $ bit $ succ s, s))
        arr :: PArray (Int,Int) Int
        arr = distribute bnds 0 add [((bit s, s), 1)] nf
        nf _ 0 = []
        nf (bs, i) val = [((setBit bs j, j), mul val $ g ! (i, j)) | j <- [0 .. pred s], not $ testBit bs j]

PArray に assocs を作るのをさぼって、Ix.range で添え字を回した。

提出結果：1543ms, 320MiB