ABC306 A～F をHaskellで

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

# A - Echo

[問題 ABC306A](https://atcoder.jp/contests/abc306/tasks/abc306_a)

シグネチャを決める。

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

## 結果

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

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

# B - Base 2

[問題 ABC306B](https://atcoder.jp/contests/abc306/tasks/abc306_b)

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

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

シグネチャを決める。

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

## 結果

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

```haskell
import Data.Word

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

# C - Centers

[問題 ABC306C](https://atcoder.jp/contests/abc306/tasks/abc306_c)

シグネチャを決める。

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

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

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

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

## 結果

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

```haskell
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](https://atcoder.jp/contests/abc306/tasks/abc306_d)

シグネチャを決める。

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

状態遷移を整理する。

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

逆に見て、

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

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

## 結果

```haskell
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](https://atcoder.jp/contests/abc306/tasks/abc306_e)

シグネチャを決める。

```haskell
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に実装した。

```haskell
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](https://atcoder.jp/contests/abc306/submissions/42515416)した。

同じアルゴリズムで、`Arrray.IO`でAの配列を管理した版は[2384 ms, 113MB](https://atcoder.jp/contests/abc306/submissions/42564516)だった。

# F - Merge Sets

[問題 ABC306F](https://atcoder.jp/contests/abc306/tasks/abc306_f)

シグネチャを決める。

```haskell
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](https://atcoder.jp/contests/abc306/submissions/42715845)
- `Data.Vector.Unboxed.Mutable`で実装したセグメント木の版で [7AC, 26TLE](https://atcoder.jp/contests/abc306/submissions/42716267) （むしろこっちのが遅い？）
- 命令型言語わからないのでJavaに上をベタ移植したつもりの版で [7AC, 16WA, 10TLE](https://atcoder.jp/contests/abc306/submissions/42739948)　（ACの数は変わっていないし、バグってWAになっているし）

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

## 公式解説版

式変形をなぞってみる。

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

```math
\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}$

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

```math
\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}
```

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

この二重シグマは、

1. スコア $p = 0$、集合 $T = \emptyset$ とする。
2. $i = N, \dots, 1$ について、順に以下を行う。
   1. $T \leftarrow T \cup S_i$ と、$S_i$を足しこむ
   1. $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には $O(\log n)$ で動く `Data.Set.lookupIndex` というものがある。集合に対して、その要素が小さい方から数えて何個めか（0始まり）を返す、まさにここの $c(x,A)$ （ただし $x \in A$）である。なので座標圧縮云々は必要ない。

$A_{i,x}$は整数なので、`Data.IntSet` を使いたくなるが、残念ながらこちらには `lookupIndex` がない。`split` (なんと$(O(1))$) で $x$ 未満の要素の集合に分割してから `size` で要素数を数えると、なぜかこれが $O(n)$ という罠がある。
（なお `Data.Set` の `size` は$O(1)$、`split` は $O(\log n)$ ）
-->
ところで、Haskellの `Data.Set` は、ある値未満と超過に集合を分割する `split` 関数があり、$O(\log n)$ で動作する。また要素数は `size` で $O(1)$ で数えられる。なので座標圧縮云々を気にせずに計算を記述できる。

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

## 結果

```haskell
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](https://atcoder.jp/contests/abc306/submissions/42761646)（制限4sec）だった。

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

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

後者ですると

```haskell
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](https://atcoder.jp/contests/abc306/submissions/42762019)になった。

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