ABC409 A～F をHaskellで

# A - Conflict

[問題 ABC409A](https://atcoder.jp/contests/abc409/tasks/abc409_a)

シグネチャを決める。

```haskell
abc409a :: Int    -- N
        -> String -- T
        -> String -- A
        -> Bool   -- 答え
```

同じ位置に `o` があれば `True` にする。

## 結果

要は `zip` を使えばいい。

```haskell
abc409a _n t a = any (('o','o') ==) $ zip t a -- 素直な表現

abc409a _n t a = elem ('o','o') $ zip t a     -- lintがこうしろと横槍を入れる

abc409a _n t a = or $ zipWith p t a           -- タプルを作らずに済ませる
  where
    p 'o' 'o' = True
    p  _   _  = False
```

# B - Citation

[問題 ABC409B](https://atcoder.jp/contests/abc409/tasks/abc409_b)

シグネチャを決める。

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

> $A$ に $x$ 以上の要素が重複を含めて $x$ 回以上現れる。

という条件を満たす最大の整数 $x$ を見つける。
$A_i$ を降順にソートしたものを $B_i$ とする。
すると、添字 $k$ の値以上の要素は $k$ 回（以上）現れている。
つまり、$i \leq B_i$ を満たす最大の $i$ が答え。

## 結果

```haskell
import Data.List

abc409b :: Int -> [Int] -> Int
abc409b n = fst . head . dropWhile (uncurry (>)) . zip [n, pred n ..] . sort
```

# C - Equilateral Triangle

[問題 ABC409C](https://atcoder.jp/contests/abc409/tasks/abc409_c)

シグネチャを決める。

```haskell
abc409c :: Int   -- N
        -> Int   -- L
        -> [Int] -- di
        -> Int   -- 答え
```

大前提として $L$ は3の倍数である必要がある。
$d_i$ を累積し $\bmod L$ した値で円周上の絶対位置 $r_i$ を作り、位置ごとに点の個数を数える。
$0 \leq a < L/3$ に対して、$L/3 \leq b < 2L/3$, $2L/3 \leq c < L$ が正三角形を作るので、その組み合わせの個数を数える。

## 結果

```haskell
import Data.Array.Unboxed
import Data.List.Split

abc409c :: Int -> Int -> [Int] -> Int
abc409c n l ds
  | r /= 0    = 0
  | otherwise = sum $ map product $ transpose $ chunksOf l3 $ elems cnt
  where
    (l3, r) = divMod l 3
    cnt = accumArray (+) 0 (0, pred l) [(a, 1) | a <- scanl add 0 ds] :: UArray Int Int
    add x y = let z = x + y in if z >= l then z - l else z -- modなしで済ませる、ケチ臭い高速化
```

せっかく配列に入れたのだから `chunksOf` など使わなくてもいい気もする。

```haskell
zipWith (*) (map (cnt !) [0 .. pred l3]) $
zipWith (*) (map (cnt !) [l3 ..])
            (map (cnt !) [l3 + l3 ..])
```

# D - String Rotation

[問題 ABC409D](https://atcoder.jp/contests/abc409/tasks/abc409_d)

シグネチャを決める。テストケース一つ分の計算をする。

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

なるべく早い位置から動かす方が辞書順で影響が大きいので、前の方から可能性を考える。
$S$ の文字を順に $C_1, C_2, \dots$ とし、
$C_i > C_{i+1}$ なら、$i = \ell$ とした結果は辞書順で小さくなる。
$C_i \leq C_{i+1}$ なら、より後方で実施するべき。

選んだ文字 $C_i$ をどこに挿入するべきか。
これも、辞書順での影響が大きくなるように、できるだけ前の方にするべき。
$C_i < C_k$ なら、$r + 1 = k$ とするべき。
$C_i \geq C_k$ なら、もっと後ろにするべき。

結局、$C_i > C_{i+1}$ な最初の文字 $C_i$ を抜いて、
それ以降の $C_i < C_k$ な最初の文字の前、そのような $C_k$ が見つからなかった場合は末尾に入れる。
都合のいい $C_i$ が見つからなかったときは、$\ell = r = 1$ とすることで元の文字列そのままが最小になる。

## 結果

```haskell
abc409d :: Int -> String -> String
abc409d _n s = loop1 s
  where
    loop1 (a:b:cs) | a > b = b : loop2 a cs
    loop1 (a:bs) = a : loop1 bs
    loop1 "" = ""

    loop2 c as@(a:_) | c < a = c : as
    loop2 c (a:as) = a : loop2 c as
    loop2 c "" = [c]
```

# E - Pair Annihilation

[問題 ABC409E](https://atcoder.jp/contests/abc409/tasks/abc409_e)

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

```haskell
abc409e :: Int     -- N
        -> [Int]   -- x_i
        -> [[Int]] -- u_i, v_i, w_i
        -> Int     -- 答え
```

ABC408F では一直線に行っていた平滑化を、木の上で行うちょっとした変種。
一般のグラフだと気が狂うがここでは木なので、葉から順に寄せて行けばできる。

注目している頂点に対して、その子孫から平滑化したお釣りの電荷が集まってくるので、
この頂点では自身の電荷に加えてそれらを合わせた分を親に向けて平滑化する。
これを流すときに、親との辺の抵抗値を掛けただけの電力が消費される。

なので、dfsでは、消費電力合計と、しわ寄せた電荷の合計の両方を集計する。

## 結果

```haskell
import Data.Array

abc409e :: Int -> [Int] -> [[Int]] -> Int
abc409e n xs uvws = ans
  where
    x = listArray (1,n) xs
    g = accumArray (flip (:)) [] (1,n) $
        [(u,(v,w)) | u:v:w:_ <-uvws] ++
        [(v,(u,w)) | u:v:w:_ <-uvws]
    (_0, ans) = dfs 0 1 0
    dfs :: Int -- 親
        -> Int -- ノード ここを電荷0にする
        -> Int -- 親からノードへの辺の抵抗値w
        -> ( Int  -- 戻ってきた電荷量
           , Int) -- 累積した答え
    dfs p v wv = (da, sum zs + wv * abs da)
      where
        (ds, zs) = unzip [dfs v c w | (c, w) <- g ! v, c /= p]
        da = sum ds - x ! v
```

# F - Connecting Points

[問題 ABC409F](https://atcoder.jp/contests/abc409/tasks/abc409_f)

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

```haskell
abc409f :: Int      -- N
        -> Int      -- Q
        -> [[Int]]  -- x_i, y_i
        -> [[Int]]  -- query_i
        -> [String] -- 答え
```

## 考える

頂点の連結状況を管理するにはUnion-Findを使えばいけるだろう。
次に連結する2点とは、距離が短くまだ連結でないものということなので、優先度付きキューで距離を管理するのだろう。
頂点が追加されるので、既存の頂点の座標も全て覚えておく必要がある。

ということで、以下の状態を考える。

```haskell
import qualified Data.Heap as PQ

type State =
  ( Int                   -- k   : 存在する点の個数
  , [[Int]]               -- xys : 点の座標Xi,Yiの逆順のリスト
  , PQ.Heap (Int,Int,Int) -- pq  : 距離、点番号ふたつの優先度付きキュー
  , UnionFind             -- uf  : 点番号から連結か判定するUnionFind
  )
```

クエリには次のように対応する。

クエリ1 `[1,a,b]` に対して、`(a,b)` を `k+1` 番めの点として追加する。
また既存の点全て `xys` との距離を測り、`pq` に投入する。

クエリ2 `[2]` に対して、
- キュー先頭のエントリから取り出し、連結なものを全て捨てる
- 連結でないエントリ `(d,i,j)` が見つかったら、答えは `d`
- キュー先頭のエントリから取り出し、距離が `d` である組を全て Union-Find で連結にする

クエリ3 `[3,u,v]` に対して、Union-Find に問いあわせた結果が答え

初期の点 N 個も、クエリ1と同じロジックで登録すればよい。

```haskell
mhd :: [Int] -> [Int] -> Int
mhd (x:y:_) (z:w:_) = abs (x - z) + abs (y - w)

abc409f _n _q xys qus = catMaybes ans
  where
    st0 = (0, [], PQ.empty, newUF)
    st1 = foldl mode1 st0 xys                     -- XiYiはクエリ1として処理
    (_st2, ans) = mapAccumL step st1 qus
    step st (1:xy)    = (mode1 st xy, Nothing)    -- クエリ1は出力なし
    step st (2:_)     = mode2 st                  -- クエリ2
    step st (3:u:v:_) = (st, Just $ mode3 st u v) -- クエリ3は状態変化なし

mode1 :: State -> [Int] -> State
mode1 (k, xys, pq, uf) xy = (k1, xy:xys, PQ.union pq pq1, uf)
  where
    k1 = succ k
    pq1 = PQ.fromList [(mhd xy zw, k1, j) | (j, zw) <- zip [k, pred k ..] xys]

mode2 :: State -> (State, Maybe String)
mode2 (k, xys, pq0, uf0) = loop1 pq0
  where
    loop1 pq =
      case PQ.uncons pq of     -- 距離は短いけど連結な対を捨てる
        Nothing -> ((k, xys, pq, uf0), Just "-1")
        Just ((d, i, j), pq1) | findUF uf0 i j -> loop1 pq1
                              | otherwise -> loop2 d pq uf0
    loop2 d pq uf =
      case PQ.uncons pq of       -- 距離 d な対を全て連結する
        Just ((e, i, j), pq1) | d == e -> loop2 d pq1 (uniteUF uf i j)
        _ -> ((k, xys, pq, uf), Just $ show d)

mode3 :: State -> Int -> Int -> String
mode3 (_,_,_,uf) u v = if findUF uf u v then "Yes" else "No"
```

immutableなUnion-Findは省略。
提出結果：[ACx28, TLEx17](https://atcoder.jp/contests/abc409/submissions/66673403)

`mode2`は `PQ.dropWhile` や `PQ.span` を使うと自前でループせず書けるが、これをやるとTLEがドンと増えてしまった。後者はともかく前者は納得いかない。

```haskell
mode2 (k, xys, pq0, uf0)
  | PQ.null pq1 = ((k, xys, pq1, uf0), Just "-1")
  | otherwise   = ((k, xys, pq3, uf1), Just $ show d)
  where
    pq1 = PQ.dropWhile (\(_,i,j) -> findUF uf0 i j) pq0 -- 距離は短いけど連結な対を捨てる
    (d, _, _) = PQ.minimum pq1
    (pq2, pq3) = PQ.span (\(e,_,_) -> d == e) pq1       -- 距離 d な対を全て連結する
    uf1 = foldl' (flip ($)) uf0 [\uf -> uniteUF uf i j | (_,i,j) <- PQ.toUnsortedList pq2]
```

## 純粋計算を捨てて

これ以上の性能を求めるには、immutableなUnion-FindとData.Heapを捨てて、 命令型配列の速度に頼らざるを得ない。

ヒープとしては、AtCoder では [`Data.Vector.Algorithms.Heap`](https://hackage-content.haskell.org/package/vector-algorithms-0.9.1.0/docs/Data-Vector-Algorithms-Heap.html)が使える。

なぜか、ベクタに格納する値に対して `Ord` 型クラスを要求する代わりに、比較関数を引数として毎回渡す形になっている。

```haskell
type Comparison e = e -> e -> Ordering
```

`compare` を渡せばいいように思うが、ひっくり返す `flip compare` 必要がある。
というのは、これは**命令型配列のヒープソート**を作ることを目標にしたヒープ機構なので、アルゴリズムの教科書に書いてあるヒープソートのやり方のとおり、配列の前半をヒープ、後半をソート済みのデータにして、ヒープの先頭に浮かび上がった最大値をヒープ末尾の値と交換し、ヒープを1マス縮小することを繰り返し、最終的に昇順ソートされた配列を作るようになっているからである（と想像される。ドキュメントには一言も説明はない。）

具体的な設計に入る。

点を追加したり、ヒープの要素を増減させたりと、配列に実際に入っている要素数が変動するので「末尾に追加する」操作が欲しい。
（C++ の `.push_back()` メソッドがいつもうらやましい。）
ベクタ自体は可変なデータ構造でも、末尾位置のインデックスを引数で別に取り回すのはスマートでないので、それをバンドルしたデータ構造を作っておく。

```haskell
import qualified Data.Vector.Unboxed.Mutable as MUV
import Data.IORef

type IOVPlus a = (IORef Int, MUV.IOVector a)

newPlus :: MUV.Unbox a => Int -> IO (IOVPlus a)
newPlus n = do
  p <- newIORef 0
  v <- MUV.new n
  return (p, v)

pushBack :: MUV.Unbox a => IOVPlus a -> a -> IO ()
pushBack (rk, v) x = do
  k <- readIORef rk
  MUV.write v k x
  writeIORef rk (succ k)

setK :: IOVPlus a -> Int -> IO ()
setK (rk, _) = writeIORef rk

getK :: IOVPlus a -> IO Int
getK (rk, _) = readIORef rk
```

`IOVPlus` を流用して、データ個数を管理してくれるヒープデータ構造として使えるようにするAPIを追加する、というか `Data.Vector.Algorithms.Heap` の薄いラッパを被せる。

```haskell
doHeapify :: (Ord a, MUV.Unbox a) => IOVPlus a -> IO ()
doHeapify (rk, v) = do
  k <- readIORef rk
  heapify (flip compare) v 0 k

pushPQ :: (Ord a, MUV.Unbox a) => IOVPlus a -> a -> IO ()
pushPQ (rk, v) x = do
  k <- readIORef rk
  heapInsert (flip compare) v 0 k x
  writeIORef rk (succ k)

nullPQ :: IOVPlus a -> IO Bool
nullPQ pq = (0 ==) <$> getK pq

peekPQ :: MUV.Unbox a => IOVPlus a -> IO a
peekPQ (_,v) = MUV.read v 0

popPQ :: (Ord a, MUV.Unbox a) => IOVPlus a -> IO ()
popPQ (rk, v) = do
  k <- pred <$> readIORef rk
  pop (flip compare) v 0 k
  writeIORef rk k
  return ()
```

上の純粋実装では $X_i, Y_i$ を登録するのにクエリ1を流用していたが、これをすると $O(n \log n)$ かかる。
最初のデータは配列に一気に突っ込んで、`heapify` で一度にヒープ化すると $O(n)$ でできるのでそのための `doHeapify` である。
またここでは $N$ 個の点間の距離 $N(N-1)/2$ 個を投入するので実際には $O(N^2 \log N)$ から $O(N^2)$ に下がる。
（$N \leq 1500$ だけど。）

あとのアルゴリズムは純粋実装と同じにして、[AC, 734ms, 114MB](https://atcoder.jp/contests/abc409/submissions/66657839) という結果になった。
<!--
（現在提出されているACしたHaskell実装では、コード量最短、実行時間最速、消費メモリ最小である。）
-->

# G - Accumulation of Wealth

公式解説によると convolution を使って計算するらしいが、そこまでに至る議論がまったくギリシャ語なので、読めるようになるまで大事にとっておく。