ABC359 A～F をHaskellで

# A - Count Takahashi

[問題 ABC359A](https://atcoder.jp/contests/abc359/tasks/abc359_a)

シグネチャを決める。

```haskell
abc359a :: Int      -- N
        -> [String] -- Si
        -> Int      -- 答え
abc359a _n = length . filter ("Takahashi" ==)
```

入力データは2種類しかないので、判定条件は `(('T' ==) . head)` で十分だったか。

# B - Couples

[問題 ABC359B](https://atcoder.jp/contests/abc359/tasks/abc359_b)

シグネチャを決める。

```haskell
abc359b :: Int  -- N
        -> Int  -- Ai
        -> Int  -- 答え
```

一人飛ばして等しいものがある位置の数を数える。

## 結果

```haskell
abc359b _n as = length $ filter id $ zipWith (==) as (drop 2 as)
```

`filter`の役割が「前の段で判断の済んだ真値を選り分ける」はさみしいので、
等しいかどうかも判定させてあげるべきなのかもしれない。

```haskell
abc359b _n as = length $ filter (uncurry (==)) $ zip as (drop 2 as)
```

これなら内包表記の方がコンパクトかもしれない。

```haskell
abc359b _n as = sum [1 | (a,c) <- zip as (drop 2 as), a == c]
abc359b _n as = length [() | (a,c) <- zip as (drop 2 as), a == c]
```

# C - Tile Distance 2

[問題 ABC359C](https://atcoder.jp/contests/abc359/tasks/abc359_c)

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

```haskell
abc359c :: [Int]  -- Sx,Sy
        -> [Int]  -- Tx,Ty
        -> Int    -- 答え
```

- 同じ高さにいる、つまりY座標が等しいとき
  - さらにそのY座標が偶数のとき、X座標を÷2した結果の差が壁の枚数
  - 奇数のときは一つずれるので、X座標をsuccしてから上をする
- 高さが違うとき、その高さの差をdyとして、高さを合わせるためにY軸方向に移動するとき、
  - まっすぐ上下に移動
  - 一つ左に移動
  - 一つ右に移動
- を任意に選べて、結局、X座標の差をdyだけ埋めることができる。
もっとも都合よく上下移動したとして、あとはX軸方向の移動について考えればよい。

## 結果

```haskell
abc359c [sx,sy] [tx,ty] = dy + abs (div sx1 2 - div tx2 2)
  where
    dy = abs $ sy - ty
    tx1 | sx < tx   = max sx (tx - dy)
        | otherwise = min sx (tx + dy)
    (sx1, tx2) = if odd sy then (succ sx, succ tx1) else (sx, tx1)
```

# D - Avoid K Palindrome

[問題 ABC359D](https://atcoder.jp/contests/abc359/tasks/abc359_d)

シグネチャを決める。

```haskell
abc359d :: Int    -- N
        -> Int    -- K
        -> String -- S
        -> Int    -- 答え
```

$K \leq 10$ がポイント。気に掛ける必要のある長さは大したことない。

Sの前 $i-1$ 文字めまで調べて、よい文字列の個数を、それぞれ、その末尾 $K-1$ 文字の状況ごとに分けて数えておく。
次の $i$ 文字めについて、手前 $K-1$ 文字にその文字を続けたパターンに対して個数を足し込む。ただし、続けた $K$ 文字が回文になる場合はしない。
`'?'` のときは両方の可能性について調べる。

AとBに2進数の0,1を割り当てることで、$2^K$ 要素の整数添え字配列で数えられる。
最初の$K-1$文字について、それより前に仮想的なAでもBでもない文字があるとする、というやり方ができないので、そこまでを前もって調べて、配列の初期値を構成する。

## 結果

Dにしてはちょっと長くなった。

```haskell
import Data.Array
import Data.Bits
import Data.List

abc359d :: Int -> Int -> String -> Int
abc359d _n k s = summ $ elems cntsN
  where
-- Kビットの2進数が回文か、の表
-- K/2ビットの整数を数え上げて、回文なビットパターンを全て生成
    isPal :: Array Int Bool
    isPal = accumArray (||) False (0, 2 ^ k - 1)
            [ (ii, True)
            | i <- [0 .. 2 ^ div (succ k) 2 - 1] :: [Int]
            , let ii = foldl' (.|.) i
                       [ bit (pred k - j)
                       | j <- [0 .. pred $ div k 2]
                       , testBit i j]
            ]
-- 前K-1文字とそれ以降
    (s1,s2) = splitAt (pred k) s
-- 前K-1文字で初期パターン
    i1s :: [Int]
    i1s = foldl' istep [0] s1
    istep is c = foldr (shiftAB c) [] is
    shiftAB c x =
      case c of
        'A' -> (xA :)
        'B' -> (xB :)
        '?' -> (xB :) . (xA :)
      where
        xA = shiftL x 1
        xB = xA .|. 1
-- カウント配列の初期値
    cnts0, cntsN :: Array Int Int
    cnts0 = accumArray (+) 0 (0, 2 ^ pred k - 1) [(i, 1) | i <- i1s]
-- K文字め以降でカウント続行
    cntsN = foldl' cstep cnts0 s2
    cstep cnts c = cnts1
      where
        cnts1 = accumArray add 0 (0, 2 ^ pred k - 1)
                [ (clearBit i1 (pred k), cnt)
                | (i, cnt) <- assocs cnts, cnt > 0
                , i1 <- shiftAB c i [], not $ isPal ! i1
                ]

modBase :: Int
modBase = 998244353
reg :: Int -> Int
reg x = mod x modBase
add :: Int -> Int -> Int
add x y = reg $ x + y
summ :: [Int] -> Int
summ = reg . sum
```

公式解説では、この解答と同様にビットでやって同様にコードが煩雑になった版と、文字列のままで実装して簡潔（かつ重い）コードでも間に合っている版と両方が示されている。文字列のままでも間に合うんだ…

## 文字列のままでする解法

やることは上と変わらないけれど、ビットに直さないので遅いけどシンプル。特に、前K-1文字について先行で処理する代わりに、AでもBでもない文字をいれておくことで決して回文にならないようにする、一種の番兵が使える。

```haskell
import qualified Data.Map as M
import Data.List

abc359d :: Int -> Int -> String -> Int
abc359d _n k s = reg $ sum $ M.elems cntsN
  where
    cnts0 = M.singleton (replicate (pred k) 'x') 1
    cntsN = foldl' step cnts0 s
    step cnts c = M.fromListWith add
        [ (take (pred k) s1, cnt)
        | (s, cnt) <- M.assocs cnts
        , s1 <- next s c
        , s1 /= reverse s1
        ]
    next s '?' = ['A' : s, 'B' : s]
    next s  c  = [c : s]
```

めっちゃ短くなった。この版は[551ms, 42MB](https://atcoder.jp/contests/abc359/submissions/54995096)、ビットでする上の版は[21ms, 17MB](https://atcoder.jp/contests/abc359/submissions/54861528)

# E - Water Tank

[問題 ABC359E](https://atcoder.jp/contests/abc359/tasks/abc359_e)

シグネチャを決める。

```haskell
abc359e :: Int   -- N
        -> [Int] -- Hi
        -> [Int] -- 答え
```

変にややこしく考えてわからなくなってしまった。
安定した状態では、水位は水源の方から階段状に低くなる一方の形状になるので、その階段のフチの位置とその高さをスタックに入れて管理する。
次の壁 $H_i$ が一番手前の壁より低ければ、幅1×高さ$H_i$だけの水を貯めて、次に溢れる。
$H_i$の方が高い場合、それより手前で$H_i$以上の高さの壁の位置までを全て$H_i$の高さまで水を貯めて、それから溢れる。$H_i$まで満ちた後は、その間にあって水没した壁は全て無視できるので、壁スタックからそれらを取り除くことができる。

## 結果

水源の位置には番兵な、高さ `maxBound` な壁を立てておく。

```haskell
abc359e :: Int -> [Int] -> [Int]
abc359e _n hs = loop 1 [(0, maxBound)] (zip [1..] hs)
  where
    loop :: Int -- 流れた総水量
         -> [(Int, Int)] -- 壁の位置と高さのスタック
         -> [(Int, Int)] -- 次に考えるべき壁
         -> [Int]  -- 答え
    loop _ _ [] = []
    loop v stk@((j,hj):_) ((i,hi):ihs)
      | hj > hi   = v1 : loop v1 ((i,hi):stk ) ihs
      | otherwise = v2 : loop v2 ((i,hi):stk2) ihs
      where
        v1 = v + hi
        (stk1, stk2) = span ((hi >=) . snd) stk
        is = map fst stk1 ++ [fst $ head stk2]
        w = sum $ zipWith (*) (map ((hi -) . snd) stk1) $ zipWith (-) is (tail is)
        v2 = v1 + w
```

# F - Tree Degree Optimization

[問題 ABC359F](https://atcoder.jp/contests/abc359/tasks/abc359_f)

シグネチャを決める。

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

さっぱりわからなかったので[フレンズさんのヒント](https://x.com/kyopro_friends/status/1804518188761440393)を見る。
> アライグマ「F問題は、「di≧1かつΣdi=2N-2」なら必ずそういう木が作れるのだ！　ということは、最初全部di=1にしておいて、「diを1増やすのにかかるコスト」が安い順に取り出せるプライオリティーキューを用意して残りのN-2回を割り当てればいいのだ！」
> フェネック「「di≧1かつΣdi=2N-2なら必ずそういう木が作れる」ことの証明は、diが大きい方から2つの頂点の間に辺を結んでNが1小さい問題に帰着すれば帰納法で言えるよ」

「そういう木」というのがわからず、理解に時間がかかった。

あるグラフが木であるならば、辺の数は$N-1$で、全体が連結である。
そのようなグラフは、全ての頂点の次数は1以上で、次数の合計は$2N-2$である。
それは割と自明。
逆に、そのような任意の次数の割り当てに対して、木が成立すること。
それをフェネックが言っているらしい。

これの証明は、公式解説にある、下から考える方がわかりやすかった。
$N=2$のとき、二つの頂点はどちらも次数1で、互いを接続して木ができて終了。
$N>2$のとき、鳩ノ巣原理により必ず次数1な頂点が存在するのでそれと、次数が1より大きい頂点も必ず存在するのでそれを接続する。次数1な方はもうこれ以上接続されないので無視すると、頂点が1減り、次数の合計が2減った問題に帰着される。

この性質を納得してしまえば、全頂点にまずは次数1を割り当て、残り$N-2$を、最もコストが安く済むように割り当てるために優先度付きキューを使うことも大体わかる。

ただここで、優先度として使うべき値が、「現在のそのノードのコスト」ではダメで「次数を1増やすことによるコストの増分」でないといけないのがよくわからない。「次数を1増やしたときのコスト」でよさそうな気もするが、例3で失敗する。

## 結果

```haskell
import qualified Data.Heap as PQ

abc359f :: Int -> [Int] -> Int
abc359f n as =
    foldl' (\acc (_,a,d) -> acc + a * d * d) 0 $ -- 総和がこたえ
    (!! (n - 2)) $                              -- N-2 本加えたら完了
    iterate step $                             -- 辺の追加を繰り返し
    PQ.fromList [(3 * a, a, 1) | a <- as]     -- キューの初期状態
  where
    step pq = PQ.insert (x + a + a, a, succ di) pq1
      where
        Just ((x, a, di), pq1) = PQ.uncons pq
```

平方数の差分は奇数 $(x + 1)^2 - x^2 = 2x + 1$ なので、初回の優先度は $3 A_i$ 以降は $+2A_i$ で、コストの増分は求められる。

最後に、キューの内容を全てリストで取り出し、$A_i, d_i$ から答えを求めるところを、リストにしてから処理するように：

```haskell
reg $ sum $ map (\(_,a,d) -> a * d * d) $ PQ.toUnsortedList $
```

と書いたら、最後のケースでTLEした。これを、`Data.Heap`自身の`Foldable`で計算するようにしたら通った。どうしてそんな差が出たのかもよくわからない。

# G - Sum of Tree Distance

公式解説に加えていっぱいユーザ解説が生えている、いろいろな解き方ができる面白い問題らしいのはわかった。