ABC340 A～F をHaskellで

# A - Arithmetic Progression

[問題 ABC340A](https://atcoder.jp/contests/abc340/tasks/abc340_a)

シグネチャを決める。

```haskell
abc340a :: Int   -- A 初項
        -> Int   -- B 末項
        -> Int   -- D 公差
        -> [Int] -- 答え
-- 高階な
abc340a a b d = takeWhile (b >=) $ iterate (d +) a
-- 内包表記
abc340a a b d = [a, a + d .. b]
```

# B - Append

[問題 ABC340B](https://atcoder.jp/contests/abc340/tasks/abc340_b)

シグネチャを決める。

```haskell
abc340b :: Int      -- Q
        -> [[Int]]  -- query i
        -> [Int]    -- 答え
```

問題文の列は末尾をトップとするスタックなので、逆向きにリストで扱う。

## 結果

```haskell
abc340b _q qs = loop [] qs

loop _ [] = []
loop xs ([1,x]:qs) = loop (x:xs) qs
loop xs ([2,k]:qs) = xs !! pred k : loop xs qs
```

# C - Divide and Divide

[問題 ABC340C](https://atcoder.jp/contests/abc340/tasks/abc340_c)

シグネチャを決める。

```haskell
abc340c :: Int  -- N
        -> Int  -- 答え
```

## 素直な解法

$x$が偶数のときは、$x/2$をふたつ書く。
$x$が奇数のときは、$\lfloor x/2 \rfloor$ と それ+1 を書く。
除算の回数が増えるが、$\lfloor x/2 \rfloor$ と $\lfloor (x+1)/2 \rfloor$ と考えてもいい。

黒板をそのままシミュレートすると大変なことになるので、それぞれの数が何個書かれているかをマップで保持し、同じ数は一度だけ、一斉に書き換える。つまり大きいものから順に処理し、コストは個数倍して数える。

```haskell
abc340c :: Int -> Int
abc340c n = loop 0 $ IM.singleton n 1

loop :: Int -> IM.IntMap Int -> Int
loop acc im
  | x < 2 = acc
  | True  = loop acc1 im2
  where
    ((x, cnt), im1) = IM.deleteFindMax im
    acc1 = acc + x * cnt
    im2 = IM.insertWith (+) (div (succ x 2) 2) cnt $
          IM.insertWith (+) (div x 2) cnt im1
```

## 公式解説のやり方

[公式解説 by Nyaan](https://atcoder.jp/contests/abc340/editorial/9246) を浅く読むと、メモ化再帰をすることで高速に計算できるとある。
というか、`@cache`とデコレータをつけると関数が勝手にメモ化されるとか、Pythonズルくない？

もう少しちゃんと読むと、黒板に現れる数が $\left \lfloor \frac{N}{2^n} \right \rfloor$ と $\left \lceil \frac{N}{2^n} \right \rceil$ だけしかないという重要な分析がある。なのでメモ化できる。上の解法はどんな値が来るかは気にしていないので無駄がある。

$N \leq 10^{17}$ 全体を配列でメモ化するのは無茶だが、メモするべき引数がごくわずかで、先に計算できるなら固定的なマップで何とかなる。$\left\lfloor \frac{N}{2^n} \right\rfloor$ と $\left\lceil \frac{N}{2^n} \right\rceil$ とは、右にシフトしていった値と、それ+1のことである。

```haskell
import Data.IntMap

abc340c :: Int -> Int
abc340c n = im ! n
  where
    im = fromList
         [ (x, f x)
         | k <- takeWhile (0 <) $ iterate (flip div 2) n
         , x <- [k, succ k]]
    f 1 = 0
    f x = x + im ! div x 2 + im ! div (succ x) 2
```

~~右シフトした値と+1しか現れないことの証明は煩雑だそうで。~~
### 追記：右シフトした値と+1しか現れないことの証明

最初の$N$から操作をした回数が同じ数をグループと考える。
操作回数$m$のグループは $\lfloor N / 2^m \rfloor$ と $+1$ だけからなること（下限は必ずあり、+1はないこともある）を証明する。

最初$N$ひとつである場合も含む、下限値だけからなる場合。
その値が偶数のとき、$[2k]$を操作すると $[k, k]$ となる。
その値が奇数のとき、$[2k+1]$を操作すると $[k, k+1]$ となる。
いずれも条件を満たす。

小さい方の値が偶数のとき、$[2k, 2k+1]$ を操作すると $[k, k, k, k+1]$ となる。これは条件を満たす。

小さい方の値が奇数のとき、$[2k+1, 2k+2]$ を操作すると $[k, k+1, k+1, k+1]$ となる。これは条件を満たしている。

wwwww (”which was what we wanted”)

<!--
どんな数になるかの予測はできないけれど、メモ化できる程度に収まるという直観が働くなら、その向きで構築する。
ってそれは普通の、メモを順次更新するメモ化再帰か、再計算で爆発するかのどちらかだ。
-->

## 兎を狩るにも全力を尽くす獅子の解法

[ユーザ解説 by MMNMM](https://atcoder.jp/contests/abc340/editorial/9247)の式を使うと、$O(1)$で答えが得られてしまう。
証明も付いているけれど、背理法による正しさの証明でなくて、どうやったらそんな式が導出されるのかという構成的な証明でないと技が盗めない。

と、そんな凄いことをしておいて、サンプルコードでは最上位ビット以下を全て塗りつぶす計算と `popCount` をベタに手書きしてみたりして、なんだかアンバランス。

写経しておく。

```haskell
import Data.Bits

abc340c :: Int -> Int
abc340c n = n * succ fn - shiftL 1 fn
  where
    fn = f $ pred n

f x = popCount . g 32 . g 16 . g 8 . g 4 . g 2 . g 1 $ x
  where
    g k x = x .|. shiftR x k
```

[フェネックさんいわく](https://twitter.com/kyopro_friends/status/1756314593184661613)
> 「この考え方を基本に考察要素を少し入れて難しくした問題がAGC044Aだから考えてみてねー」

### 追記：ある程度構成的な証明

<details><summary>与式の証明</summary>

解説にも書いてあるけれど、問題のプロセスに沿って何が起きているのかという観点から説明を試みる。

与式：$f(N) = \lceil \log_2(N) \rceil$として 答え $C(N) = N(f(N) + 1)-2^{f(N)}$

$N$が2のべきであるかどうかで場合分けして考える。

- $N = 2^K$ つまり2のべきであるとき

$f(N) = K$
$N$を2で割っていくと、黒板の様子は順に $2^K$ が $2^0$ 個、$2^{K-1}$ が $2^1$ 個、…、$2^1$ が $2^{K-1}$ 個、$2^0 = 1$ が $2^K$ 個、と推移する。
一つのステップで支払うコストは$N$円なので、$K$ステップの総額は$NK$円。
$C(N) = N(K+1) - 2^K = NK + N - N = NK$ なので一致する。

- $2^{K-1} < N < 2^K$ つまり2のべきでないとき

$f(N) = K = J + 1$ とおく。
黒板の様子を同様に考える。$i$ステップでの小さい方の値の個数を$a_i$、+1な値の個数を$b_i$とする。
初期状態ステップ0では$N$が1つなので $a_0=1, \\, b_0=0$である。
ステップ$i$での小さい方の値が偶数のとき、上で考えたように、次のステップでは$a_{i+1} = 2a_i + b_i, \\, b_{i+1} = b_i$ である。
一方奇数のときは $a_{i+1} = a_i, \\, b_{i+1} = 2a_i + b_i$ である。
ここで、どちらにせよ $a_{i+1} + b_{i+1} = 2(a_i + b_i)$ である。また $a_0 + b_0 = 1$ とあわせて $a_i + b_i = 2^i$ である。

小さい方の値とは $\left\lfloor N/2^i \right\rfloor$ で、それが奇数とは$N$の第$i$ビットが1であること、偶数とはビットが0であることに対応する。

改めて $b_0=0$である。
$N$の第$i$ビットが0のとき、$b_{i+1}=b_i$である。
$N$の第$i$ビットが1のとき、$b_{i+1}=a_i + 2b_i = b_i + (a_i + b_i) = b_i + 2^i$ である。
これは、第$i$ステップでは、$N$の第$i$ビットを$b_i$に写していることに他ならない。
$b_J$は$N$の最上位ビットの手前までを写した値となり、$b_J = N - 2^J$ と表せる。

黒板に戻る。$J$ステップ後の状態は 1 が $a_J$ 個、2 が $b_J$ 個である。ここまでにかかったコストは$NJ$で、さらに次のステップで $b_J$個の 2 を倍の個数の1に書き換えて完了する。
総コストは $NJ + 2(N - 2^J)$ である。

与式を振り返ると
$C(N) = N(K+1) - 2^K = N(J+2) - 2^(J+1) = NJ + 2N - 2 \cdot 2^J = NJ + 2(N - 2^J)$
で、めでたく一致した。

wwwww (”which was what we wanted”)

</details>

# D - Super Takahashi Bros.

[問題 ABC340D](https://atcoder.jp/contests/abc340/tasks/abc340_d)

シグネチャを決める。$A_i,B_i,X_i$は手抜きする。

```haskell
abc340d :: Int      -- N
        -> [[Int]]  -- Ai,Bi,Xi
        -> Int      -- 答え
```

一見、前から考えて、各ステージは一度だけ考えれば済むようにも見える。素直に順にステージを辿っていく進み方を下限として、$X_i$で飛ぶのは前に進む場合だけ気にすればいい、と。
しかし、素直に進むやり方の途中にめちゃめちゃ重いステージが入っていて、その前後を一度引き返す形でまたぐ経路が最短、という意地悪なケースもありうるので、素直にダイクストラ法で計算する。

## 結果

```haskell
abc340d :: Int -> [[Int]] -> Int
abc340d n abxs = runST $
  do
    dist <- dijkstra n (abxA !) 1
    readArray dist n
  where
    abxA = listArray (1,n) $
           [[(i1, a), (x, b)] | (i1, a:b:x:_) <- zip [2..] abxs] ++
           [[]] -- ステージNの分

dijkstra :: Int                  -- 頂点数N (1～N)
         -> (Int -> [(Int,Int)]) -- 隣接頂点とその辺の重み、グラフの情報
         -> Int                  -- 開始点
         -> ST s (STArray s Int Int)
dijkstra n edges start = -- (略)
```

[全容は提出で。](https://atcoder.jp/contests/abc340/submissions/50193367)

ここでは命令型配列として`Data.Array.ST`を使った。
命令型配列を使うアルゴリズムをSTモナドで作ると、純粋な計算の中に戻すことができる。
ただし、複数の純粋な計算の間を渡り歩くことは結局できない。
そうしたいときはIOモナドで作る必要があるが、今度は純粋な計算には戻せない。（`unsafePerformIO`はナシで）
どっちもどっちで悩ましい。

# E - Mancala 2

[問題 ABC340E](https://atcoder.jp/contests/abc340/tasks/abc340_e)

シグネチャを決める。

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

箱の中のボールの個数を配列 $X[0 \sim N-1]$ で持ち、添え字は $\bmod N$ で巡回するとして、
一つの操作は、

- $C = X[B_i]$ とする
- $X[B_i] \leftarrow 0$ とする
- $X[B_i + 1 \sim B_i + C]$ に +1 する（$C>N$のとき、2以上増加する要素もある）

となる。
配列に対して区間の更新を行う必要があるので、贅沢に区間更新セグメント木を配列代わりに使って、状況を素直にシミュレーションすることができる。

一点注意として、$C$は$N$よりもずっと大きな値になりうるので、配列を何周も回る可能性があり、それを+1ずつするのではなく、増加する値を一気に足し込む必要がある。

- $B_i + C < N$ で回り込まないとき：$B_i + 1$から$B_i + C$に1ずつ足し込む
- 回り込むとき、$(B_i + C) \div N = q \dots r$ として、$B_i+1$から$N-1$までに1足し込み、0から$N-1$までに$q - 1$足し込み、1から$r$までに1足し込む。
なおここで、$(B_i + 1) \bmod N$とやるのは間違いなので注意。$q=1$のときは全体への足し込みは省略できる。

## 純粋計算での区間更新配列の実装

完全な区間更新セグメント木を使う代わりに、区間更新できる配列を実装してみる。つまり、区間に対するクエリはできない。

配列の添え字は0始まりとする。
保持するべき情報は、重ねる演算の操作とその単位元、葉の枚数、木そのもの、からなる。

```haskell
-- Span Update Array
data SUA a = SUA (a->a->a) a Int (SUAT a)
```

木は、自分の子に重ねる値と、左右の子を持つ。
葉の場合は子を`undefined`にしておく。現在位置が葉かどうかは、再帰的に降りていくときには判るので情報を残しておく必要はない。

```haskell
data SUAT a = Node a (SUAT a) (SUAT a)
```

初期値リストとその長さから配列を作る。要素数は2のべきになるように補われる。

```haskell
fromListSUA :: (a->a->a) -- 重ね演算
            -> a      -- 単位元
            -> Int    -- 要素数N
            -> [a]    -- 要素0～N-1
            -> SUA a  -- 区間更新配列
fromListSUA op u n xs = SUA op u w t
  where
    w = head . dropWhile (n >) $ iterate (2 *) 1
    t = loop w xs
    loop 1 []    = Node u undefined undefined
    loop 1 (x:_) = Node x undefined undefined
    loop w xs    = Node u (loop w2 xs) (loop w2 (drop w2 xs))
      where w2 = div w 2
```

添え字を指定して要素を読み出す。

```haskell
readSUA :: SUA a -- 配列
        -> Int   -- 添え字
        -> a     -- 読み出した値
readSUA (SUA op u w t) k = loop 0 w u t
  where
    loop _p 1 acc (Node x _ _) = op acc x
    loop  p w acc (Node x l r)
      | k < m = loop p w2 acc1 l
      | True  = loop m w2 acc1 r
      where
        w2 = div w 2
        m  = p + w2
        acc1 = op acc x
```

区間と値を指定して、値を重ねる。

```haskell
pileSUA :: Int   -- 区間下限（含む）
        -> Int   -- 区間上限（含まない）
        -> a     -- 重ねる値
        -> SUA a -- 配列
        -> SUA a -- 更新された配列
pileSUA a b v (SUA op u w t) = SUA op u w $ loop 0 w t
  where
    loop p w t@(Node x l r)
      | b <= p || q <= a = t
      | a <= p && q <= b = Node (op v x) l r
      | otherwise = Node x (loop p w2 l) (loop (p + w2) w2 r)
      where
        q = p + w
        w2 = div w 2
```

先頭からN要素をリストで読み出す。一つずつ`readSUA`するより速いはず。

```haskell
toListSUA :: Int   -- 要素数
          -> SUA a -- 配列
          -> [a]   -- 結果
toListSUA n (SUA op u w t) = loop 0 w u t []
  where
    loop p 1 acc (Node x _ _) rest
      | n <= p = rest
      | True   = op acc x : rest
    loop  p w acc (Node x l r) rest
      | n <= p = rest
      | True    = loop p w2 acc1 l $ loop (p + w2) w2 acc1 r rest
      where
        w2 = div w 2
        acc1 = op acc x
```

どれも、二分木の左右に再帰的に計算を進め、注目している区間と現在のノードの守備範囲とで場合分けをする構造は似たようなコードである。（のでボイラープレートとしてまとめられるとかっこいいかもだが…）

以上で構築した配列データ型を使って、次のように解くことができる。

```haskell
abc340e :: Int -> Int -> [Int] -> [Int] -> [Int]
abc340e n _m as bs = toListSUA n $ foldl step (fromListSUA (+) 0 n as) bs
  where
    step :: SUA Int -> Int -> SUA Int
    step sua b =
      case divMod (b1 + pred c) n of
        (0, r) -> pileSUA b1 (succ r) 1 sua1
        (1, r) -> pileSUA b1 n 1 $ pileSUA 0 (succ r) 1 sua1
        (q, r) -> pileSUA b1 n 1 $ pileSUA 0 (succ r) 1 $ pileSUA 0 n (pred q) sua1
      where
        c = readSUA sua b
        b1 = mod (succ b) n
        sua1 = pileSUA b (succ b) (negate c) sua
```

この実装は[無事TLEした。](https://atcoder.jp/contests/abc340/submissions/50195108)

通常のセグメント木の方法で、命令型配列に二分木を埋め込んでやったら[1216ms, 212MB](https://atcoder.jp/contests/abc340/submissions/50196093)でACした。

## ユーザ解説の方法

[すこし定数倍の軽い方法 by MMNMM](https://atcoder.jp/contests/abc340/editorial/9267)
問題の内容を素直に受け取ると、ゾルトラーク化した区間更新セグメント木を使うだけ問題に見えるが、工夫することで区間更新のないセグメント木でも解ける、という話。

箱$i$のボールの個数を配列に格納する代わりに、配列の$0$から$i$までの累積和が箱$i$のボールの個数になるように情報を維持する。すると、次のように操作を読み替えられる。

- 配列要素の読み出しの代わりに、累積和を求める
- 区間更新の代わりに、累積和の変化するエッジだけを操作する
- 一点更新は（上と同様に）区間更新で行う

これらは、普通のセグメント木、またはフェニック木 (Binary Indexed Tree) で効率的に実現できる。

なるほど方針は理解できたので実装まではしないで済ませる。

# F - S = 1

エフ引くエス イコール いち、と読んでしまった。

[問題 ABC340F](https://atcoder.jp/contests/abc340/tasks/abc340_f)

シグネチャを決める。

```haskell
abc340f :: Int   -- X
        -> Int   -- Y
        -> [Int] -- 答え ([A,B]または[-1])
```

$(0,0)$と$(X,Y)$を通る直線を $L : Yx - Xy = 0$ とする。
$(0,0)$と$(X,Y)$の距離を $R = \sqrt{X^2+Y^2}$ とおく。
$L$と$(A,B)$との距離を $H$ とすると、考える三角形の面積は底辺×高さ÷2で $S = RH/2$ である。
この$H$は高校数学の問題で、点と直線の距離公式 https://manabitimes.jp/math/857 より $H = |YA - XB|/R$
すると $S = |YA - XB|/2$ これが1になるとは $YA -XB = \pm 2$
（ある $a,b$ で $Ya - Xb = -2$とできたとすると、$Y(-a) - X(-b) = -(-2) = 2$ なので、マイナスは無視できる。）

さて、$YA -XB = 2$ を満たす整数解 $(A,B)$ を見つけるのも、（少し逸脱した）高校数学の問題で、一次不定方程式ax+by=cの整数解 https://manabitimes.jp/math/674 より、
拡張ユークリッドの互除法で $YA - XB = G (= \gcd(Y,-X))$ の$G$に加えてこの等式を満たす一つの解$A_0,B_0$が求まるので、2が$G$の倍数ならば、$A_0,B_0$を$2/G$倍すると$YA -XB = 2$の解の一つになる。2が$G$の倍数でないとき、整数解はない。

## 結果

```haskell
abc340f x y
  | abs g > 2 = [-1]
  | otherwise = [a * g2, b * g2]
  where
    ((a,b),g) = extendGCD y (negate x)
    g2 = div 2 g

extendGCD :: Integral b => b -> b -> ((b, b), b)
extendGCD a b = loop a b
  where
    loop a 0 = ((1, 0), a)
    loop a b = let ((y, x), d) = loop b (mod a b)
               in  ((x, y - div a b * x), d)
```

$(A,B)$から$L$と平行に移動した先に、条件を満たす格子点は無限に存在する。
問題によっては、それらの中から条件を満たすものを探し出す必要があるかもしれないが、この問題はとても条件が緩いので、拡張ユークリッドの互除法で最初に見つけた解そのままで十分。

# G - Leaf Color

[アライさんいわく](https://twitter.com/kyopro_friends/status/1756318113350729738)
> 木DPパートだけでF問題くらい難しいからじっくり考えるのだ！

今日のところはこのくらいにしといてやるのだ！（降参です。）