ABC392をHaskellで

# A - Shuffled Equation

[問題 ABC392A](https://atcoder.jp/contests/abc392/tasks/abc392_a)

シグネチャを決める。

```haskell
abc392a :: Int -> Int -> Int -- A1,A2,A3
        -> Bool              -- 答え
abc392a x y z = x * y == z || y * z == x || z * x == y
```

掛け算に順序はないので、$B_3$ となるのがどれかの3通りを全て試すのが話が早いだろう。

# B - Who is Missing?

[問題 ABC392B](https://atcoder.jp/contests/abc392/tasks/abc392_b)

シグネチャを決める。

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

問題文の集合的な言い回しのせいで `IntSet` を使いそうになったが、$N \leq 1000$ 程度なので配列で済ませた方が早い。

## 結果

```haskell
import Data.Array

abc392b :: Int -> Int -> [Int] -> [Int]
abc392b n _m as = [i | (i,True) <- assocs arr]
  where
    arr = accumArray (&&) True (1,n) [(a,False) | a <- as]
```

$O(NM)$ にはなるが、何ならすごいH本1章の範囲だけでも解ける。

```haskell
abc392b n _m as = [i | i <- [1 .. n], notElem i as]
```

# C - Bib

[問題 ABC392C](https://atcoder.jp/contests/abc392/tasks/abc392_c)

シグネチャを決める。

```haskell
abc392c :: Int   -- N
        -> [Int] -- Pi
        -> [Int] -- Qi
        -> [Int] -- 答え
```

- $Q_i$ から $i$ を逆引きする配列で「$j$ （独立だけど、ここは問題文の変数を変えてほしかったな）が書かれたゼッケンを着けている人の番号 $k$」がわかる。
- 「番号 $k$ の人が見つめている人の番号 $\ell$」は、$i$ から $P_i$ をひく配列で取り出せる。
- 「番号 $\ell$ の人が着けているゼッケン番号 $m$」は、$i$ から $Q_i$ をひく配列で取り出せる。

## 結果

```haskell
import Data.Array

abc392c :: Int -> [Int] -> [Int] -> [Int]
abc392c n ps qs = [i2q ! (i2p ! (q2i ! j)) | j <- [1 .. n]]
  where
    i2p = listArray (1,n) ps
    i2q = listArray (1,n) qs
    q2i = array (1,n) $ zip qs [1 ..]
```

# D - Doubles

[問題 ABC392D](https://atcoder.jp/contests/abc392/tasks/abc392_d)

シグネチャを決める。サイコロの情報は横着する。
（$K_i$ は捨てても `length` で取り返せるけど、絶対に必要なので勿体ないし）

```haskell
abc392d :: Int     -- N
        -> [[Int]] -- Ki, Aij
        -> Double  -- 答え
```

$N \leq 100$ なので総当たりでやればよいのだろう。
目の数の範囲 $A_{i,j} \leq 10^5$ が大きいのと、同じ目が複数の面にある可能性から、配列でも `IntSet` でもなく `IntMap Int` で、目の数に対してその目のある面の個数を数えておく。

あとは総当たりで、サイコロ$i$と$j$について、共通する目の個数の積の和をとり、$K_i K_j$ で割ればゾロ目の確率になる。

## 結果

```haskell
import qualified Data.IntMap as IM

abc392d :: Int -> [[Int]] -> Double
abc392d n kass = maximum
    [ fromIntegral nume / fromIntegral deno
    | (k1, im1):d2s <- tails ds, (k2, im2) <- d2s
    , let nume = sum (IM.elems $ IM.intersectionWith (*) im1 im2)
    , let deno = k1 * k2 ]
  where
    ds = [(k, IM.fromListWith (+) [(a, 1) | a <- as]) | k:as <- kass]
```

# E - Cables and Servers

[問題 ABC392E](https://atcoder.jp/contests/abc392/tasks/abc392_e)

シグネチャを決める。

```haskell
abc392e :: Int -- N
        -> Int -- M
        -> [[Int]] -- Ai, Bi
        -> [[Int]] -- 答え
```

Union-Find を使って、連結にするのに貢献する最小限の辺とそうでない辺を仕分けする。（クラスカル法）
使わずに余った辺を集めておく。
Union-Find の結果の、分割の代表元のリストを取り出す。
このリストの先頭から順に、余った辺を使ってさらに連結にしていく。

余った辺とは、両辺が同じ分割に繋がっているということ。
なので、つなぎ替えるのはどちらの辺でも構わない。それを、違う分割の代表元に繋ぐ。
具体的には、代表元リストの先頭を r として、自分の接続先が r と異なるなら、r に繋ぐ。
自分の接続先が r なら、リストの次の要素と繋ぐ。
このときどちらにせよ、r が代表元でなくなるように Union-Find で次の連結を行う。（それが制御できない Union-Find の場合は、代表元でなくなった方をリストから除く。）

代表元リストが長さ1になったとき、全体が連結になったので完了する。

## 結果

```haskell
abc392e :: Int -> Int -> [[Int]] -> [[Int]]
abc392e n m uvs = runST $
  do
    imroot <- newArray (1, n) True :: ST s (STUArray s Int Bool)
    uf <- newUF (1, n)
-- 最小スパニング木に必要な辺を使い、使わなかった辺を集める
    iuvs <- foldM (\iuvs iuv@(i, u:v:_) -> do
      r <- uniteUF uf u v
      case r of
        Nothing     ->                              return (iuv:iuvs) -- 使わなかった辺をとっておく
        Just (a, _) -> writeArray imroot a False >> return iuvs       -- 辺を使ったので代表が一人減る
      ) [] (zip [1 ..] uvs)
-- 代表を集める
    roots <- map fst . filter snd <$> getAssocs imroot
-- roots が singleton になるまで繰り返すことの、
-- iuvsを一つずつ見て、その両端の代表と異なる代表に接続しなおす、を行う
-- その結果、ufとrootsが同期して変化する
    ans <- loop uf [] iuvs roots
    return $ [length ans] : ans
  where
-- rootsが1つになったら完了
    loop _uf ans _iuvs [_] = return ans
-- (u,v) は代表元 r1 の分割にはいっている。r1と別の分割の代表元とを繋ぐ。
-- r1 ＝ r のときは unite r (head rs) として r を消す これはuまたはvをhead rsにつなぎ替えたことに相当する
-- r1 ≠ r のときは unite r r1 として r を消す これはu or vをrにつなぎ替えたことに相当する
    loop uf ans ((i,u:v:_):iuvs) (r:rs) = do
      r1 <- getRoot uf u -- vも同じ
      let (s, w) = if r1 == r then (head rs, head rs) else (r1, r)
      uniteUF uf r s
      loop uf ([i, u, w]:ans) iuvs rs
```

# F - Insert

[問題 ABC392F](https://atcoder.jp/contests/abc392/tasks/abc392_f)

シグネチャを決める。

```haskell
abc392f :: Int   -- N
        -> [Int] -- Pi
        -> [Int] -- 答え
```

時計を巻き戻して考えて、$P_i$ の位置に $i$ を入れるということは、それ以前の位置 $p$ で $P_i \leq p$ なものは、実際には $p+1$ を指している、というズレを累積して追跡していけば、$N$ から 1 までを指定された位置に収めるだけの仕事になる。

## 直接解法1

`Data.Set.elemAt` は、順序集合の $i$ 番目の要素を $O(\log N)$ で取り出せる凄いやつ。
$\\{1, \dots, N\\}$ から始めて、集合の $P_i$ 番目の要素が $i$ の入る位置となり、この位置はもう使用済みなので集合から除いて続ける、とすれば解ける。

```haskell
import qualified Data.Set as S
import Data.Array.Unboxed

abc392f :: Int -> [Int] -> [Int]
abc392f n ps = elems arr
  where
    arr :: UArray Int Int
    arr = array (1, n) $ snd $ mapAccumR step (n, S.fromList [1 .. n]) ps
    step (i, s) p = ((pred i, S.delete x s), (x, i))
      where
        x = S.elemAt (pred p) s
```

定数倍がきつくて[1758 ms](https://atcoder.jp/contests/abc392/submissions/62587734)かかったけれど、間に合った。

## 時計の巻き戻しもしない、どストレート解法

子の要素数をメモした二分木を使えば、先頭から $i$ 番目の位置に要素を挿入する操作は、木の高さを適切に制御すれば $O(\log N)$ でできるはず。
削除はなしで挿入だけが行われるなら、ある程度いい加減なアルゴリズムでも何とかならないだろうか。

二分木は、葉にだけデータがあり、それらを繋ぐ分岐点には葉の枚数を持たせておく。
部分木の葉の枚数を取り出す補助関数も立てておく。

```haskell
data Tree a = Leaf a | Node Int (Tree a) (Tree a)

size :: Tree a -> Int
size (Leaf _) = 1
size (Node s _ _) = s
```

気が早いが、完成した木から、左から順に葉を取り出してリストに並べる関数

```haskell
elemsT t = recur t []
  where
    recur (Leaf x) rest = x : rest
    recur (Node _ lt rt) rest = recur lt $ recur rt rest
```

左から i 番目（0始まり）の位置に要素 x を挿入するとき、対象が葉であれば、位置は0か1しかない。

```haskell
-- 仮
insertT 0 x t@(Leaf _) = Node 2 (Leaf x) t
insertT 1 x t@(Leaf _) = Node 2 t (Leaf x)
```

対象が分岐のとき、左の部分木の中にその位置があるなら、左に分岐する。
右の部分木の中にあるなら、右に分岐する。
ちょうど中間のとき、左の部分木の末尾につけても、右の部分木の先頭につけても、どちらでも正しいので、このとき、小さい方に付けることにする。

```haskell
-- 仮
insertT k x (Node s lt rt)
  | k < ls -> Node (succ s) (insertT k x lt) rt        -- 左へ
  | ls < k -> Node (succ s) lt (insertT (k - ls) x rt) -- 右へ
-- k == ls のとき
  | ls <= size rt -> ... -- 左へ
  | otherwise     -> ... -- 右へ
  where
    ls = size lt
```

再帰的に木を降りていくのはいいが、再帰呼び出しの結果を木で受け取り、反対側は何もせず接続するのはバランスに配慮がなく、$O(N)$ な木を作らされてしまう。
例えば「左へ」の再帰呼び出しの結果 `insertT k x lt = Node s ll mm` という木ができ、これが `rt` と比べて巨大になりうる。
ここには3つの部分木 `ll`, `mm`, `rt` があり、その順序を入れ替えるわけにはいかないので、簡単にできることと言えば、どちらを先に `Node` で結ぶか、`Node * (Node * ll mm) rt` とするか `Node * ll (Node * mm rt)` とするかを選ぶぐらいである。
3つの部分木の大きさを考えて、`rt`が最大なら前者を、`ll`が最大なら後者を選ぶとマシになる。さて、`mm` が最大のときはどうすればよいか？

再帰呼び出しの結果が木として返る、と考えると、ここで行き詰まる。再帰呼び出しの結果を、ここでいう二つの部分木 `ll`, `mm` の対（またはリスト）で返すと考えると、`mm` が最大のとき、最大ということは大きさは1より大きく、つまり分岐点であることが確定しているので、それを `Node * ml mr` とすると、`[Node * ll ml, Node * mr rt]` と二つに割くことで、マシな結果になると期待できる。

以上をまとめると次のような挿入関数ができる。

```haskell
insertAtT i x t = Node (size lt + size rt) lt rt -- 最後は二つをまとめて終わり
  where
    lt:rt:_ = recur i t
    recur 0 t@(Leaf _) = [Leaf x, t] -- 葉への挿入は自明
    recur 1 t@(Leaf _) = [t, Leaf x]
    recur k (Node s lt rt) =
      case compare k ls of
        LT -> goLeft
        GT -> goRight
        EQ | ls + ls <= s -> goLeft
           | otherwise    -> goRight
      where
        ls = size lt
        goLeft  = post $ recur k lt ++ [rt]
        goRight = post $ lt : recur (k - ls) rt
    post ts@[t1, t2, t3]                         -- 3つを2つにまとめなおす後処理
      | smax == s1 = [t1, Node (s2 + s3) t2 t3]
      | smax == s3 = [Node (s1 + s2) t1 t2, t3]
      | otherwise  = [Node (s1 + size t2l) t1 t2l, Node (size t2r + s3) t2r t3]
      where
        ss@[s1,s2,s3] = map size ts
        smax = maximum ss
        Node _ t2l t2r = t2 -- t2が最大なら、Leafということはない
```

これができたら、指示を前から順に実行すれば解ける。

```haskell
abc392f :: Int -> [Int] -> [Int]
abc392f n (_1:ps) = elemsT $ foldl' step (Leaf 1) $ zip ps [2 ..]
  where
    step t (pi, i) = insertAtT (pred pi) i t
```

タイムは [1441 ms](https://atcoder.jp/contests/abc392/submissions/62595311) と、`Data.Set` 版を上回った。

## 想定解

上の直接解法1で、「現在$P_i$番目と呼べるものは最終的にどの位置にくるものだったか」を `Data.Set` に問いあわせているが、これをセグメント木で行う。
配列要素の初期値を全て1にしておくと、先頭からの合計が$P_N$となる位置が、最終ステップでNを書き込むべき位置。使った枠の要素を0にしていく。
手順が進んで、現在の歯抜けな状況で、先頭からの合計が$P_j$となる位置は、値$j$が最終段階で落ち着く位置になる。
先頭からの総和がちょうど$P_j$となる位置、を求めるのには、「セグメント木というデータ構造に対して二分探索を行う」のではなく、セグメント木の内部構造を利用して直接二分探索を行う。

通常のセグメント木で、区間 [a,b) に関する問いあわせは
- 現在注目している部分木が
  - 完全に区間に包含されているなら、ノードに記された値を返す
  - 共通部分が全くないなら、ゼロ値を返す
  - 部分的に交差しているなら、左右の部分木に同じ問い合わせをし、結果を演算でまとめて結果とする

という再帰計算になる。
述語 p を満たすような区間 [0, x) の x の最大値を二分探索で見つけるには、
- 現在注目している区間 [a,b) に対応する部分木と、[0, a) の値 v について
　- 左部分木の根に記された値 w について p(v + x) が真なら、a,bの中点を m として、区間 [m, b) に対応する右部分木と [0, m) の値 v + x について再帰する
　- さもなくば、[a, m) に対応する左部分木と v について再帰する
- 注目している部分木が葉になったら、現在位置を x として返す

という再帰計算になる。

結果：[363ms, 83MB](https://atcoder.jp/contests/abc392/submissions/62618464)
直接解法1と、やっていることに大差はないはずなのだが、効率は全然違った。

# G - Fine Triplets

[問題 ABC392G](https://atcoder.jp/contests/abc392/tasks/abc392_g)

シグネチャを決める。

```haskell
abc392g :: Int   -- N
        -> [Int] -- Si
        -> Int   -- 答え
```

なんもわからん。

https://x.com/kyopro_friends/status/1888225425236861197

FFTが出てくるんだ。
実数に対するFFTなら学部の実験でPC-9801相手にプログラムを書いたような気がするが、それっきりだ。

[NTT(数論変換)を使って2つの多項式の積を求めるプログラム](https://maehrm.hatenablog.com/entry/2024/05/09/200649)にあるプログラムが、リスト操作に素直に対応付けることかできたので移植してみた。
結果：[TLEx32](https://atcoder.jp/contests/abc392/submissions/62673977)

[フーリエ変換を用いて多項式の掛け算を行う](https://blog.applibot.co.jp/2020/12/16/fourier-transform/)のRustコードをよくよく読むと、毎回の計算で配列のほぼ全ての要素を更新するので、immutable arrayでも書けるとわかってそうしてみた。
結果：[AC, 2831ms, 369MB](https://atcoder.jp/contests/abc392/submissions/62745146) どうにか間に合った。

mutable array にすればもっと速くなるはずとやってみた（[STUArray](https://atcoder.jp/contests/abc392/submissions/62829720), [Vector.Unboxed.Mutable](https://atcoder.jp/contests/abc392/submissions/62830005)）が、だめだった。何がいけないのだろうか。