ABC243 A～G をHaskellで

E, Gについて、公式解説とは少し異なる別解を見つけた。

# A - Shampoo

[問題 ABC243A](https://atcoder.jp/contests/abc243/tasks/abc243_a)

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

```haskell
abc243a :: [Int] -- V,A,B,C
        -> Char  -- 答え
```

## ゆるいやり方

順に調べてもA問題なら間に合う。

```haskell
abc243a [v,a,b,c] = loop v
  where
    loop v0
      | v0 < a = 'F'
      | v1 < b = 'M'
      | v2 < c = 'T'
      | otherwise = loop $ v2 - c
      where
        v1 = v0 - a
        v2 = v1 - b
```

## かしこいやり方

グルグル回らなくても、$A+B+C$ で割った余りが最後の周回の初期値になる。

```haskell
abc243a [v,a,b,c]
  | v0 < a     = 'F'
  | v0 - a < b = 'M'
  | otherwise  = 'T'
  where
    v0 = mod v (a + b + c)
```

# B - Hit and Blow

[問題 ABC243B](https://atcoder.jp/contests/abc243/tasks/abc243_b)

シグネチャを決める。

```haskell
abc243b :: Int   -- N
        -> [Int] -- Ai
        -> [Int] -- Bi
        -> (Int,Int) -- 答え
```

$A, B$ 中に重複はないので、$A$ と $B$ の共通の要素の個数から一つめの答えを引けば、二つめの答えになる。

## 結果

```haskell
import Data.List

abc243b :: Int -> [Int] -> [Int] -> (Int, Int)
abc243b n as bs = (cnt1, cnt2 - cnt1)
  where
    cnt1 = length $ filter id $ zipWith (==) as bs
    cnt2 = length $ intersect as bs
```

`Data.IntSet` を用いてもっと効率的にすることもできる。

# C - Collision 2

[問題 ABC243C](https://atcoder.jp/contests/abc243/tasks/abc243_c)

シグネチャを決める。

```haskell
abc243c :: Int     -- N
        -> [[Int]] -- Xi,Yi
        -> String  -- S
        -> Bool    -- 答え
```

左右にしか動かないので、Y座標の値ごとに分けて、
- 右に移動する最も左から出発する人
- 左に移動する最も右から出発する人

のX座標を集め、同じ高さでかつ向かい合っている人がいるか調べればよい。

## 結果

```haskell
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

abc243c :: Int -> [[Int]] -> String -> Bool
abc243c _n xys s = or [r < l | (y,l) <- IM.assocs ls, Just r <- [IM.lookup y rs]]
  where
    ls = IM.fromListWith max [(y, x) | ('L', x:y:_) <- zip s xys]
    rs = IM.fromListWith min [(y, x) | ('R', x:y:_) <- zip s xys]
```

# D - Moves on Binary Tree

元々、この問題が[AtCoder NoviSteps](https://atcoder-novisteps.vercel.app/)から
スタックの問題として参照されていてたどり着いた。

[問題 ABC243D](https://atcoder.jp/contests/abc243/tasks/abc243_d)

シグネチャを決める。

```haskell
abc243d :: Int     -- N
        -> Int     -- X
        -> String  -- S
        -> Int     -- 答え
```

現在位置を二進数の桁ごとの（逆順の）リストとして保持する。

## 結果

```haskell
import Data.List

abc243d :: Int -> Int -> String -> Int
abc243d _n x s = deBits $ foldl' step (enBits x) s
  where
    step bs 'U' = tail bs
    step bs 'L' = 0 : bs
    step bs 'R' = 1 : bs

enBits :: Int -> [Int]
enBits 0 = []
enBits x = r : enBits q
  where
    (q,r) = divMod x 2

deBits :: [Int] -> Int
deBits = foldr step 0
  where
    step d acc = acc + acc + d
```

# E - Edge Deletion

[問題 ABC243E](https://atcoder.jp/contests/abc243/tasks/abc243_e)

シグネチャを決める。

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

## 考える

$N \leq 300$ というヒントから、よく訓練された競技プログラマはワーシャルフロイド法を呼び出して、
全ての頂点間の距離を計算する。頂点 $i$ から $j$ までの距離を $d[i,j]$ とする。
ワーシャルフロイド法の計算量は $O(N^3)$ である。

さて、$A_i, B_i, C_i$ が不要なのは、これを使わずにそれ以下の距離で行ける別ルートがあるとき。
$d[A_i,B_i] \leq C_i$
だが待ってほしい。$<$ なら確かに迂回路があるが、$=$ のときは実はこの辺を辿った距離なのかもしれない。

### 公式解説のやり方

公式ではここで「何らかの中間点 $v$ を経由した距離で以下の距離を達成できるか」で判定するようにして、
その最短距離が迂回路でも達成可能かを判断している。
$\exists v; d[A_i,v] + d[v,B_i] \leq C_i$
全ての辺について別々に、全ての頂点を中間点として試すので、
この方法の計算量は $O(MN) \leq O(N^3)$ である。

### 自分のやり方

何が欲しいかというと、ワーシャルフロイド法で計算した距離が、
複数の辺を辿ったものか、単独の辺の長さなのかを見分けたい。
そこで、距離の数値に細工を加えて、辺の長さと、それを足し合わせたものを区別できるようにする。
なお、実際の長さが等しいときは、単独よりも繋いだものの方が短いとみなされるようにする。
これにはいくつかの方法が考えられる。

- 単独辺の長さかそうでないかを1ビット追加する。
`Bool` をタプルで追加したり、長さを倍にしてビット0を使ったりできる。
足し算と比較の差し替えは、ワーシャルフロイド法のライブラリを直接触る他、
`Num` 型クラスのインスタンスにする方法もあるだろう。
- 辺の長さ `Int` と、使っている辺の本数 `Int` のタプルにする。
後者を負の数で数えることで、大小関係は期待通りになる。
- 上の、使っている本数の情報を長さの下位ビットにねじ込む
本数は 1 から M が区別できればよいので、長さ $C$ の辺を $C (M+1) - 1$ に映す。
二つの辺の和は $C_1 (M+1) - 1 + C_2 (M+1) - 1 = (C_1 + C_2)(M+1) - 2$ となり、
以降、本数の分だけ、本来の長さ $(M+1) \sum C_i$ より少ない値になる。
減らされる数の上限は $M$ なので、$M$を足してから$M+1$で割れば、本来の長さの真値は必ず復元できる。

### 結果

自分のやり方は後半が $O(M) \leq O(N^2)$ で済むので、公式解説のやり方より速い。
ワーシャルフロイド法の計算量 $O(N^3)$ の方が支配的だけど。

```haskell
abc243e :: Int -> Int -> [[Int]] -> Int
abc243e n m abcs = length [() | a:b:c:_ <- abcs, d ! (a,b) < cf c]
  where
    d = runSTUArray $ warshallFloyd (1,n) $ concat
        [[(a,b,c1),(b,a,c1)] | a:b:c:_ <- abcs, let c1 = cf c]
    cf c = pred $ c * succ m

-- warshallFloyd は省略
```

[提出コード, AC, 171ms](https://atcoder.jp/contests/abc243/submissions/65132871)

# F - Lottery

[問題 ABC243F](https://atcoder.jp/contests/abc243/tasks/abc243_f)

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

```haskell
abc243f :: [Int]   -- N,M,K
        -> [Int]   -- Wi
        -> Int     -- 答え
```

順番を考えずに、$K$ 回の抽選で当たった内訳が、賞品 $i$ が $C_i$ 回 となったとする。
$\sum C_i = K$ である。
考えたいのは賞品が $M$ 種類ということなので、$C_i$ の中で0でないものが $M$ 個な場合である。
そのような $K$ の配分は、再帰的に、もしくはそのDP化により計算できる。
残りの抽選回数 $p$ 残りの当てるべき賞品の種類数 $q$ という状況で
賞品 $i$ 以下の当選数のリスト $[C_i, C_{i+1}, \dots, C_K]$ の集合を $C[i,p,q]$ とすると、

$C[i,0,0] = \\{[C_i = 0, \dots, C_N = 0]\\}$  
$C[i,p,0] = \emptyset$ （$p > 0$ のとき）
$C[i,0,q] = \emptyset$ （$q > 0$ のとき）
$C[i,p,q] = C[i+1,p,q] \cup \\{ [C_i = c, cs] \ |\  c \in \\{1, \dots, p-q+1\\}, cs \in C[i+1, p-c, q-1]\\}$

と定義できる。
問題で考慮するべき全ての場合は $C[1,K,M]$ で得られる。

当選数リスト $[C_1, \dots, C_N]$ となるような確率は $\prod (\frac{W_i}{S})^{C_i}$ （ただし $S = \sum W_i$）である。
また、このような個数になる全ての場合は、重複のある順列組み合わせの場合の数で $K! / \prod C_i !$ である。

列を作ってから全ての列についてこれらを計算する代わりに、
$C[i,p,q]$ の計算過程にこれらの計算を融合させてプログラムにする必要がある。

## 結果

```haskell
import Data.Array

abc243f :: [Int] -> [Int] -> Int
abc243f [n,m,k] ws = pr ! (1,k,m)
  where
-- 階乗の逆数
    rf = listArray (0, k) $ scanl' (\v i -> mul v $ modRecip i) 1 [1 ..]
-- 分母 ΣWi の逆数
    de = modRecip $ sum ws
-- Wi/ΣWi 一度のくじでiを得る確率
    pA = listArray (1, n) $ map (mul de) ws
-- DP
    bnds = ((1,0,0),(succ n,k,m))
    pr = listArray bnds $ map pf $ range bnds
    pf (_,0,0) = prodd [1 .. k]
    pf (_,_,0) = 0
    pf (_,0,_) = 0
    pf (i,_,_) | n < i = 0
    pf (i,p,q) = summ
      [ prodd [rf ! c, pp, pr ! (succ i, p - c, if c == 0 then q else pred q)]
      | (c, pp) <- zip [0 .. p - pred q] $ iterate (mul (pA ! i)) 1 ]

-- モジュロ演算のコードは省略
```

[提出コード, AC, 31ms](https://atcoder.jp/contests/abc243/submissions/65137836)

# G - Sqrt

シグネチャを決める。一つのテストケースに対応する部分だけにする。

```haskell
abc243f :: Int  -- X
        -> Int  -- 答え
```

公式の「$O(X^{1/2})$ 解法」に相当するところまでは自力で見つけられたが、その先がわからず解説を見た。
解説の中の謎の計算 $dp[x] = \sum_{i=1}^{\lfloor x^{1/4} \rfloor} \left (\lfloor x \rfloor - i^2 + 1 \right) dp[i]$
を理解する代わりに、自分の解法の続きを考え直したら答えまで至ることができた。

## 考える

### 平方根の段階

$A_1 = X$ から始めて、作ることができる数列の種類数を $C(X)$ とする。
次の数は $\sqrt{X}$ 以下であり、$\sqrt 1 = 1$ と1に到達するまでどんどん小さくなる数列なので、  
$C(1) = 1$  
$C(X) = \sum_{i=1}^{\lfloor \sqrt X \rfloor} C(i)$  
となる。
$C$の累積和 $D(X) = \sum_{i=1}^X C(i)$ とすると $C(X) = D(\lfloor \sqrt X \rfloor)$ であり、  
$D(1) = 1$  
$D(X) = D(X-1) + C(X) = D(X-1) + D(\lfloor \sqrt X \rfloor)$  
と、$D(1)$ ～ $D(\lfloor \sqrt X \rfloor)$ だけを配列上に事前計算しておくことで $C(X)$ を求めることができる。

しかし $X \leq 9\times10^{18}$ から、Dの要素数は $3\times10^9$ も必要で、時間も空間も足りない。

### 4乗根の段階

$D(X)$ の漸化式を完全に展開すると  
$D(1) = 1$  
$D(X) = \sum_{i=1}^X D(\lfloor \sqrt X \rfloor)$  
つまり  
$D(X) = D(\lfloor \sqrt 1 \rfloor) + D(\lfloor \sqrt 2 \rfloor) + \dots + D(\lfloor \sqrt X \rfloor)$  
となる。しかしこれは $\lfloor \sqrt X \rfloor$ 項ではなく $X$ 項ある。

整数平方根の値を観察すると

|X|1|2|3|4|5|6|7|8|9|10|…|16|…|25|…|
|:--:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
|$\lfloor \sqrt X\rfloor$|1|1|1|2|2|2|2|2|3|3|←|4|←|5|←|

という感じに、Xが平方数になるたびに1ずつ増えて、同じ値が続く $k^2 \leq X < (k+1)^2$ の区間の要素数は $2k+1$ である。

つまり、$k = \lfloor \sqrt X \rfloor$ とすると、$k^2 \leq X < (k+1)^2$ であって、  
$D(X) = D(1) \times 3 + D(3) \times 5 + \dots + D(k) \times m$  
$= \sum_{i=1}^{k-1} (2i+1) D(i) + mD(k)$ （ここで $m = X - k^2 + 1$）

よって、最初から通しで考えると、床関数は省略して  
$C(X) = D(X^{1/2})$ （$X \leq 9 \times 10^{18}$）  
$D(Y) = \sum_{i=1}^{k-1} (2i+1) D(i) + mD(k)$ （$k = Y^{1/2}, m = Y - k^2 + 1, Y \leq 3 \times 10^9$）  
$D$と $2i+1$ を乗じた累積和は $k \leq X^{1/4} = \sqrt{3 \times 10^9} \fallingdotseq 54772.\cdots$ まで持っておけばよい。

### 結果

$D$ (`dee`) に$2i+1$を乗じた累積和を$E$ (`eee`) としている。
これは、公式と同じ計算量「前計算 $O(X^{1/4})$ クエリ $O(1)$」をもつ。

<!--
これは、公式の「前計算 $O(X^{1/4})$ クエリ $O(1)$」解法に他ならない。
$A_3 = i$ と置く云々の議論を理解しないまま、同じ場所にたどり着いてしまった。
似た要素もあるけどどうも対応付かないので、別解と理解しよう。
-->

```haskell
import Data.Array

abc243g :: Int -> Int
abc243g x = eee ! pred z + (succ y - z * z) * dee ! z
  where
    y = iSqrt x  -- X^1/2
    z = iSqrt y  -- X^1/4

ub :: Int
ub = 9 * 10^18

dee, eee :: Array Int Int
(dee, eee) = (dee, eee)
  where
    su = iSqrt ub
    qu = iSqrt su
    dee = listArray (1, qu) deelist
    deelist :: [Int]
    deelist = scanl (+) 1 [dee ! iSqrt i | i <- [2 ..]]
    eee :: Array Int Int
    eee = listArray (0, qu) eeelist
    eeelist :: [Int]
    eeelist = scanl (+) 0 $ zipWith (*) (iterate (2 +) 3) deelist

iSqrt :: Int -> Int
iSqrt n = foldr step 0 $ take 32 $ iterate (2 *) 1
  where
    step b c = let bc = b + c in if bc <= div n bc then bc else c
```

# Ex - Builder Takahashi (Enhanced version)

[アライグマ「Ex問題は、SとGを分断すればいいのだ！　アライさんは実装してないのだ……」](https://x.com/kyopro_friends/status/1502645391413903361)

撤収！