D. 簡単なDP
E. 多重集合

A - Echo

問題 ABC306A

シグネチャを決める。

abc306a :: Int     -- N
        -> String  -- S
        -> String  -- 答え

結果

Stringは[Char]なので適当に。

abc306a _n s = foldr (\c s -> c:c:s) "" s

B - Base 2

問題 ABC306B

1ビットをタグに使うために長い整数が64ビットに満たない言語に配慮して、64ビット整数を使い切る出題は避ける、という了解があったように思うんだが。MSBも使っているから符号なしでないといけないし。

実際Haskellも普通に Int だとあふれる。符号なし64ビット整数は Data.Word.Word64 にある。

シグネチャを決める。

abc306b :: [Int]   -- Ai
        -> Word64  -- 答え

結果

あとは素朴に計算するだけ。

import Data.Word

abc306b :: [Int] -> Word64
abc306b as = sum [b | (b, 1) <- zip (iterate (2 *) 1) as]

C - Centers

問題 ABC306C

シグネチャを決める。

abc306c :: Int    -- N
        -> [Int]  -- Ai
        -> [Int]  -- 答え

$A_i$ を前から見ていき、

その値が初出なら、左にあるものなので、出力しないが、一度出たことを覚えておく
二度めなら、中央なので、出力し、二度出たことを覚えておく
三度めなら、右のものなので、何もしない

四度めはありえないので、「（奇数回出たことを）覚えておく」はカウンタでなく1ビットで足りる。

結果

1ビットのカウンタは Data.IntSet で代用する。

import qualified Data.IntSet as IS

abc306c :: Int -> [Int] -> [Int]
abc306c _n as = loop IS.empty as
  where
    loop _ [] = []
    loop is (a:as)
      | IS.member a is = a : loop (IS.delete a is) as
      | otherwise      =     loop (IS.insert a is) as

D - Poisonous Full-Course

なかなかに不謹慎なネタである。

問題 ABC306D

シグネチャを決める。

abc306d :: Int          -- N
        -> [(Int,Int)]  -- Xi Yi
        -> Int          -- 答え

状態遷移を整理する。

現状態	食べない	毒消し入りを食べる	毒入りを食べる
健康	健康	健康	毒を受ける
どく	どく	健康	しに

逆に見て、

毒消し入りの料理だったとき、
- 健康な状態へは、健康なときに食べないでいる、健康なときに食べる、どく状態で食べる、で遷移する
- どく状態へは、どく状態で食べない、で遷移する
毒入りの料理だったとき、
- 健康な状態へは、健康なときに食べないでいる、で遷移する
- どく状態へは、健康なときに食べる、どく状態で食べない、で遷移する

2つの状態について、スコアの最大値を維持するDPで解ける。

結果

abc306c :: Int -> [(Int,Int)] -> Int
abc306c _n = uncurry max . foldl step (0,0)

-- 毒消し入り
step (healthy, poisoned) (0, y) = (maximum [healthy, healthy + y, poisoned + y], poisoned)
-- 毒入り
step (healthy, poisoned) (1, y) = (healthy, max (healthy + y) poisoned)

初期状態で「毒を受けてスコア0」は本来ありえないが、最初の料理が毒のとき食べたら死ぬ、最初の料理が毒消しなら食べたら「健康でスコア0」からの遷移と合流する。食べなかった場合も、健康な初期状態から食べなかった場合より不利な状況のままなのは変わらないので、これでよい。

初期状態が「お腹の痛い高橋くんがレストランにやってきた」だったら、ちゃんと考える必要がある。

E - Best Performances

問題 ABC306E

シグネチャを決める。

abc306e :: Int          -- N
        -> Int          -- K
        -> Int          -- Q
        -> [(Int,Int)]  -- Xi Yi
        -> [Int]        -- 答え

話を整理する。

大きさ$N$の書き換えられる配列$A$がある。初期値はオール0
$i = 1,\dots,Q$ について、順に以下を行う
- $(X_i, Y_i)$ の指示で、$A[X_i] \leftarrow Y_i$と更新する
- その後、$A$の大きい方$K$個の和を報告する

大きい方$K$個の値の多重集合$U$と、残りの値の多重集合$L$を追跡する。IntMapで実装すると、最大値や最小値が取り出しやすい。
また、$U$の合計も追跡する。

初期状態は $U$：0が$K$個、$L$：0が$N-K$個、となる。ここから容易に見えるように、$U$の最小値$U_{\min}$と$L$の最大値$L_{\max}$は重なることがある。

$X_i$ に対して、$A[X_i]$をそれが所属する$U$または$L$から除き、$Y_i$ に対して、$U$の要素が$K$個になるように注意して、どちらかに挿入する。この操作を最も手際よく達成しようとすると、$X_i, Y_i, U_{\min}, L_{\max}$ の大小関係により、場合分けをどうしていいかわからなくなった。

そこで、実行時間で不利になるが、単純明快になるように実装する。

まず$X_i$を除く。$U$から除いた場合、$L_{\max}$をひとつ$L$から$U$に移動させ、$U$は$K$個ある状態を維持する。
次に$Y_i$を加える。$U$に加えた場合、$U_{\min}$をひとつ$U$から$L$に移動させ、$U$は$K$個ある状態を維持する。

毎回、最後に$A[X_i]$を$Y_i$に更新する。これはmutable arrayでしないと性能上不利になる。

例外的な場合として、$K = N$ のとき、全ての値が常に$U$に属する。$L$が空なため、$L_{\max}$を触る上の方針をそのまま適用することはできないので、特別対応をする。といってもはるかに単純である。

結果

$A$の配列もIntMapでimmutableに実装した。

import Data.List
import qualified Data.IntMap as IM

abc306e :: Int -> Int -> Int -> [(Int,Int)] -> [Int]
abc306e n k q xys
  | k < n = snd $ mapAccumL stepA iniA xys
  | True  = snd $ mapAccumL stepB iniB xys
  where
    a0 = IM.empty
    iniA = (a0, IM.singleton 0 (n-k), IM.singleton 0 k, 0)
    iniB = (a0, 0)

-- Aのアクセサ
getA :: IM.IntMap Int -> Int -> Int
getA a x = IM.findWithDefault 0 x a

-- N == K の特別な場合
-- 状態は (配列A, Aの合計)
type StateB = (IM.IntMap Int, Int)

stepB :: StateB -> (Int,Int) -> (StateB, Int)
stepB (a, s) (x, y) = ((s1, a1), a1)
  where
    a1 = IM.insert x y a
    s1 = s - getA a x + y

-- 一般の場合
-- 状態は(配列A, L, U, Uの合計)
type StateA = (IM.IntMap Int, IM.IntMap Int, IM.IntMap Int, Int)

stepA :: StateA -> (Int,Int) -> (StateA, Int)
stepA (a, l, u, s) (x, y) = ((a1, l2, u2, s2), s2)
  where
    a1 = IM.insert x y a
-- ステップ1 : A[Xi]=Zを除く
    z = getA a x -- 指名された、消える値
    (lmax,_) = IM.findMax l -- Lの最大値
    (l1, u1, s1)
      | z <= lmax = (decr z l, u, s) -- ZがL側にあるとき、それを抜くだけ
      | otherwise = (decr lmax l, incr lmax $ decr z u, s - z + lmax) -- zがU側にあるとき、それを抜いた分、Lの最大値をUへ移動させる
-- ステップ2 : Yを追加する
    (umin,_) = IM.findMin u1 -- Uの最小値
    (l2, u2, s2)
      | y <= umin = (incr y l1, u1, s1) -- YがL側に入るとき、入れるだけ
      | otherwise = (incr umin l1, incr y $ decr umin u1, s1 - umin + y) -- YがU側に入るとき、UminがひとつLに押し出される

-- 多重集合に値を一つ追加
incr x im = IM.insertWith (+) x 1 im
-- 多重集合から値を一つ抜く
decr x im = IM.update f x im
  where
    f 1 = Nothing
    f k = Just $ pred k

このimmutableな実装でも5065 ms, 295MBでACした。

同じアルゴリズムで、Arrray.IOでAの配列を管理した版は2384 ms, 113MBだった。

F - Merge Sets

問題 ABC306F

シグネチャを決める。

abc306f :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- Aij
        -> Int      -- 答え

これも話がややこしい。整理する。

要素数$M$の数の集合$S_i$が$N$個あり、互いに疎である。
$S_i, S_j$ を選んだとき、両方合わせた中で $S_i$の要素の小さい順での順位の合計を $f(S_i, S_j)$ と定義する。
全ての $i < j$ な組み合わせでの $f(S_i, S_j)$ の総和を求めよ。

（うまくいかなかった）自分の考え

$A_{i,x}$は互いに異なるので、$NM$個をひとまとめにしても、どれがどの $S_i$ に属しているかはわかる。
これ全体を昇順にソートして、前から順に見ていく。
$S_i$ の要素 $A_{i,x}$ が来たとする。
このとき、スコアには、$i < j \leq N$ な $S_j$ に対する順位を足しこむ。

全ての $S_i$ それぞれについて、その要素がこれまでに何個来たかを配列$C[i]$で数えておくと、$S_i \cup S_j$ における $A_{i,x}$ の順位は $C[i] + C[j]$ である。
つまり、今回増えるスコアは
$\displaystyle \sum_{j = i+1}^N (C[i] + C[j]) = (N-i)C[i] + \sum_{j=i+1}^N C[j]$
配列のある連続区間の和を得たり、更新したりするには、セグメント木やフェニック木がよい。

セグメント木の要素数が$N$なので、更新や区間和の問い合わせ一回のコストが$O(\log N)$で、$A_{i,x}$は全部で $NM$ 個あるので、トータルで $O(NM\log N)$ となる。
$N \leq 10^4, M \leq 10^2$ なので充分間に合う、はずだった。

代数的データ型でimmutableに実装したセグメント木の版で 19AC, 14TLE
Data.Vector.Unboxed.Mutableで実装したセグメント木の版で 7AC, 26TLE （むしろこっちのが遅い？）
命令型言語わからないのでJavaに上をベタ移植したつもりの版で 7AC, 16WA, 10TLE　（ACの数は変わっていないし、バグってWAになっているし）

踏んだり蹴ったりで終了。

公式解説版

式変形をなぞってみる。

集合$S$の$x$以下な要素の個数は $|\{a \in S \ |\ a \leq x \}|$ と数学的な記法で書けるが、これを$c(x,S)$としよう。

\begin{array}{ll}
\displaystyle \sum_{1\leq i < j \leq N} f(S_i,S_j) & 求めたいもの \\
\displaystyle = \sum_{1\leq i < j \leq N} \sum_{k=1}^M c(A_{i,k},S_i \cup S_j) & fの定義 \\
\displaystyle = \sum_{1\leq i < j \leq N} \left (\sum_{k=1}^M c(A_{i,k},S_i) + \sum_{k=1}^M c(A_{i,k},S_j) \right ) & cの性質
\end{array}

ここで、内側一つめについて
$\displaystyle \sum_{k=1}^M c(A_{i,k},S_i) = \sum_{k=1}^M k = \frac{M(M+1)}{2}$
全ての $S_i$ の要素数が同じなためこれは $i,j$ に依存せず、
$\displaystyle \sum_{1\leq i < j \leq N} \frac{M(M+1)}{2} = \frac{(N-1)N}{2} \cdot \frac{M(M+1)}{2}$

二つめに関して計算を再開する。

\begin{array}{ll}
\displaystyle \sum_{1\leq i < j \leq N} \sum_{k=1}^M c(A_{i,k},S_j) & 求めたい残り \\
\displaystyle = \sum_{i = 1}^{N-1} \sum_{j=i+1}^N \sum_{k=1}^M c(A_{i,k},S_j) & Σを手抜きせず書く \\
\displaystyle = \sum_{i = 1}^{N-1} \sum_{k=1}^M \sum_{j=i+1}^N c(A_{i,k},S_j) & 2,3個めのΣを入れ替える \\
\displaystyle = \sum_{i = 1}^{N-1} \sum_{k=1}^M \Big (c(A_{i,k},S_{i+1}) + \dots + c(A_{i,k},S_N) \Big ) & 3個めのΣを展開 \\
\displaystyle = \sum_{i = 1}^{N-1} \sum_{k=1}^M c(A_{i,k},S_{i+1} \cup \dots \cup S_N) & cの性質（逆向き）
\end{array}

な、なるほど…嘘は言ってない…

この二重シグマは、

スコア $p = 0$、集合 $T = \emptyset$ とする。
$i = N, \dots, 1$ について、順に以下を行う。
1. $T \leftarrow T \cup S_i$ と、$S_i$を足しこむ
2. $S_i$の全ての要素 $A_{i,x}$ について、$c(A_{i,x}, T)$ をスコアに足しこむ。
  $p \leftarrow p + \sum_{k=1}^M c(A_{i,x}, T)$

で計算できる。

公式解説では、Aを座標圧縮してフェニック木を使って数えることができ、計算量は $O(NM \log NM)$ とある。
座標圧縮とは、$A_{i,j}$を全体で整列したときの番号に置き換えて、その番号をフェニック木の添え字として使うということ。要素数$NM$なのでそういうことになる。

ところで、Haskellの Data.Set は、ある値未満と超過に集合を分割する split 関数があり、$O(\log n)$ で動作する。また要素数は size で $O(1)$ で数えられる。なので座標圧縮云々を気にせずに計算を記述できる。

$A_{i,x}$は整数なので Data.IntSet を使いたくなるが、なぜかこちらは size が $O(n)$ という罠がある。

結果

import qualified Data.Set as S

abc306f :: Int -> Int -> [[Int]] -> Int
abc306f n m ass = div (m * succ m) 2 * div (n * pred n) 2 + p
  where
    (_, p) = foldr step (S.empty, 0) ass
    step as (s, p) = (s1, p1)
      where
        s1 = S.union s (S.fromList as)
        p1 = p + sum [S.size lower | a <- as, let (lower,_) = S.split a s]

ただし結果はギリギリの3941ms（制限4sec）だった。

やはり、計算量としては $O(\log n)$でも、分割した集合を実際に構築するのは時間がかかるだろうから、もう一つの凄い関数 lookupIndex を使うことを考える。これは、集合に含まれている要素について、それが小さい方から何番目かを返す $O(\log n)$ の関数。要素である必要があり、しかし探そうとしている a はまだ s には入っていないので、

まず a 未満の s 中の最大要素を lookupLE で探して、それに対して lookupIndex する
s1 に対して lookupIndex する。$A_i$ の分増えすぎるので、後で補正する

後者ですると

abc306f n m ass = div (m * succ m) 2 * div (n * pred n) 2 + p
                - n * div (m * succ m) 2                        -- 補正
  where
    ...
        p1 = p + sum [succ j | a <- as, let Just j = S.lookupIndex a s1]
    ...

最終的に2847msになった。

（といっても、なんかAtCoderの判定環境が重くて速度も不安定？攻撃が続いているの？）