ABC342をHaskellで

# A - Yay!

[問題 ABC342A](https://atcoder.jp/contests/abc342/tasks/abc342_a)

シグネチャを決める。

```haskell
abc342a :: String -- S
        -> Int    -- 答え
```

先頭2文字が同じなら、それは「全て同じ文字」側の字。
異なるとき、3文字めと異なる方が答えの文字。

## 結果

```haskell
abc342a s@(c1:c2:c3:_)
  | c1 == c2  = head [i | (i, c) <- zip [1..] s, c /= c1]
  | c1 == c3  = 2
  | otherwise = 1
```

# B - Which is ahead?

[問題 ABC342B](https://atcoder.jp/contests/abc342/tasks/abc342_b)

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

```haskell
abc342b :: Int     -- N
        -> [Int]   -- Pi
        -> Int     -- Q
        -> [[Int]] -- Ai,Bi
        -> [Int]   -- 答え
```

「人$P_i$がいるのは前から何番目か」を引ける配列を作っておけばよい。

## 結果

```haskell
import Data.Array

abc342b :: Int -> [Int] -> Int -> [[Int]] -> [Int]
abc342b n ps _q abs = map f abs
  where
    ia = array (1,n) $ zip ps [1..]
    f (a:b:_) = if ia ! a < ia ! b then a else b
```

# C - Many Replacement

[問題 ABC342C](https://atcoder.jp/contests/abc342/tasks/abc342_c)

シグネチャを決める。文字数が多いので `Data.ByteString` を使う。$c_i, d_i$は手抜きする。

```haskell
import qualified Data.ByteString.Char8 as BS

abc342c :: Int             -- N
        -> BS.ByteString   -- S
        -> Int             -- Q
        -> [BS.ByteString] -- ci,di
        -> BS.ByteString   -- 答え
```

## 素朴な解法

「文字 $x$ は文字 $y$ に置き換える」向きの対応を持つ配列を、$c_i \\; d_i$ について全体をスキャンして、$y$が$c_i$である要素を全て$d_i$に書き換えればよい。26要素を舐める必要があるがそれは定数。
いちいち `Array` を構築するのも無駄なので、最終的に `listArray` で配列を作るための、置き換え先要素のリストで処理する。

```haskell
import Data.Array
import Control.DeepSeq

abc342c :: Int -> BS.ByteString -> Int -> [BS.ByteString] -> BS.ByteString
abc342c n s _q cds = BS.map (aQ !) $ BS.take n s
  where
    repl c d x = if x == c then d else x
    lQ = foldl' step ['a'..'z'] cds
    step l bs = force $ map (repl (BS.index bs 0) (BS.index bs 2)) l
    aQ = listArray ('a','z') lQ
```

`BS.take n` という一見無駄なことをしているのは、手元の計算機で実行すると`BS.getLine`で取り込んだ文字列の末尾に`\r`が混入するため。Windowsのせい？

素朴にやると[1734msでギリギリ](https://atcoder.jp/contests/abc342/submissions/50701268)だけど、`force`とか`foldl'`とか小技を効かせると[152msで動く。](https://atcoder.jp/contests/abc342/submissions/50701334)

## 逆写像を用いる解法

「文字$y$に置き換えるのは文字$a,b,c,\dots$」向きの対応を持つ配列を、$c_i \\; d_i$ に対して、$c_i$ の分を全て $d_i$ に渡すという更新を行う。最後にこの配列を舐めて対応を正向きに戻す。

文字の集合はビット配列で表せるサイズなのでそうして稼ぐ。

```haskell
import Data.Bits
import Data.Array.Unboxed

abc342c :: Int -> BS.ByteString -> Int -> [BS.ByteString] -> BS.ByteString
abc342c n s _q cds = BS.map (c2d !) $ BS.take n s
  where
    d2c :: UArray Char Int  -- dになる文字集合をビット表現で
    d2c = foldl' step (listArray ('a','z') $ iterate (2 *) 1) cds
    step arr cd
      | c == d    = arr
      | otherwise = arr // [(c, 0), (d, arr ! d .|. arr ! c)]
      where
        c = BS.index cd 0
        d = BS.index cd 2
    c2d :: UArray Char Char -- d2cの逆引き
    c2d = array ('a','z')
      [ (c, d)
      | (d,cs) <- assocs d2c
      , (c,i) <- zip ['a'..'z'] [0..]
      , testBit cs i
      ]
```

immutable arrayでの実装にもかかわらず、[62msを達成した。](https://atcoder.jp/contests/abc342/submissions/50701622)

同じアルゴリズムをつたないC++で書いてみた結果は[22ms](https://atcoder.jp/contests/abc342/submissions/50713322)。

<!--
## ユーザ解説の方法

[より高速な解法 by MMNMM](https://atcoder.jp/contests/abc342/editorial/9357)は、Union-Find風の木構造で、元々とある文字だった位置に入れるべき文字が取り出せる木を作り、文字の置き換えを木のつなぎ替えにしている感じ？
やっていることは理解できたつもりなのだけど、きれいに抽象化できなくて気持ち悪い。[27ms](https://atcoder.jp/contests/abc342/submissions/50713350)
-->

# D - Square Pair

[問題 ABC342D](https://atcoder.jp/contests/abc342/tasks/abc342_d)

シグネチャを決める。

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

（自分の解法が色々と中途半端だったので、公式解説、ユーザ解説を見てリファインしたやり方を説明する。）

0は平方数である。
正の平方数を素因数分解すると、それぞれの素因数についてその個数は全て偶数になる。またそのような数だけが正の平方数である。

0でない二つの$A_i, A_j$を掛けた値が平方数になるとは、同じ素因数の和が偶数になることと等しい。
そうなるためには、両者の素因数の個数が、同じ素因数について偶奇が同じになればよい。
よって、各$A_i$について奇数個な素因数の昇順のリストを作り、それが一致する相手が、掛け合わせて平方数になる数である。
ここで、このリストはそのまま使わなくとも、掛け合わせた単一の整数で代用できる。（フレンズさんのユーザ解説を見るまでこれに気づかなかった。）この数を特徴数と呼ぼう。

次に、平方数になる組み合わせを数え上げることを考える。
$A_j$を前から順に調べていく。

$A_j$が0でないとき、特徴数が一致する、より手前の$A_i$の個数が、正の平方数を得る組み合わせの個数である。これを数えるために、特徴数をキー、その特徴数を持つ$A_i$の個数を値とする`IntMap`を維持し、取り出した個数を足し込めばよい。

また、0であるような$A_i$と掛けると、任意の$A_j$は0という平方数を得る。よって、これまでに出現した0の個数も足し込む。

一方$A_j$が0であるとき、その手前の全ての$A_i$との組み合わせで0が得られるので、$j-1$を足し込む。

0について違うやり方として、0の総数$z$だけ別で数えると、0でない数は$n - z$個で、0と0でない数の組み合わせが$n(n-z)$個、0どうしの組み合わせが$z(z-1)/2$個になる。

## 結果

前者の方法をとる。
`step`の状態の`im`は状態数をキー、個数を値とするIntMap、`zc`は出現した0の個数。
リストで読み込む`j`は今回の$A_i$より前にある数の個数（つまり$i-1$）、`a`は$A_i$そのもの。

```haskell
import qualified Data.IntMap as IM

abc342d :: Int -> [Int] -> Int
abc342d n as = sum $ snd $ mapAccumL step (IM.empty, 0) $ zip [0..] as
  where
    step (im, zc) (j, 0) = ((im, succ zc), j)
    step (im, zc) (_, a) = ((IM.insertWith (+) c 1 im, zc), zc + IM.findWithDefault 0 c im)
      where
        c = product $ singlify $ primeFactors a -- 特徴数

-- 連続する2要素が等しいとき消す
singlify :: Eq a => [a] -> [a]
singlify (x:y:xs) | x == y = singlify xs
singlify (x:xs) = x : singlify xs
singlify [] = []

--  素因数分解
primeFactors :: Int -> [Int]
primeFactors n = loop n primes
  where
    primes = 2 : 3 : [y | x <- [5,11..], y <- [x, x + 2]]
    loop n pps@(p:ps)
      | n == 1    = []
      | n < p * p = [n]
      | r == 0    = p : loop q pps
      | otherwise = loop n ps
      where
        (q,r) = divMod n p
```

## 公式解説のやり方

特徴数の作り方だけが異なる、思いきったやり方。

素因数分解すらサボって、$A_i$の最大値までの平方数を全て作っておくと、$A_i$を割りきる最大の平方数で割った商が$A_i$の特徴数になる。

上のコードとの差分だけ示す。`singlify`も`primeFactors`も不要。

```haskell
abc342d :: Int -> [Int] -> Int
abc342d n as = ...
  where
    sqs = reverse $ takeWhile (maximum as >=) $ map (^2) [1..] -- 平方数リスト
    ...
    step (im, zc) (_, a) = ...
      where
        c = head [q | sq <- sqs, (q,0) <- [divMod a sq]] -- 特徴数
```

# E - Last Train

[問題 ABC342E](https://atcoder.jp/contests/abc342/tasks/abc342_e)

シグネチャを決める。

```haskell
abc342e :: Int     -- N
        -> Int     -- M
        -> [[Int]] -- li, di, ki, ci, Ai, Bi
        -> [Maybe Int] -- 答え、到達不能のときNothing
```

駅$N$からグラフ構造の幅優先探索で調べると、経路は長いが辺の和は短い正解でなく、長い一辺を使った経路を選んでしまうのでいけない。

それぞれの駅の条件を満たす最も遅い出発時刻をキューに入れて、その降順に、つまり遅い方から、その便に間に合うさらに前の駅からの便を探す、を繰り返すことで解ける。これはダイクストラ法と同じ方針になるが、辺の長さが固定でない点で異なる。

与えられた時刻$t$と駅$B$に関して、列車群 $(l, d, k, c, A, B)$ の中で、時刻$t$以前に到着できるような最も遅い便の出発時刻は、
$l + jd + c \leq t \\; (0 \leq j < k)$ より $j = \min (k-1, \lfloor \frac{t-l-c}{d} \rfloor)$ となり、
$0 \leq j$のとき $l + jd$ 、さもなくばそのような便はない。

最初はキューに$(\infty, N)$を入れて始める。
キューの先頭の要素$(t, B)$に対して、$B$が既出ならスルーする。
初出なら、$B$の終電は時刻$t$と決める。さらに、$B$に$t$以前に到着できる便の出発時刻と出発駅$(l + jd, A)$を全てキューに追加する。
この手順をキューが空になるまで繰り返す。

## 結果

降順のキューは符号の反転で実現した。

```haskell
import qualified Data.Heap as PQ
import Data.Array
import qualified Data.IntMap as IM

abc342e :: Int -> Int -> [[Int]] -> [Maybe Int]
abc342e n m ldkcabs = [IM.lookup x ans | x <- [1..pred n]]
  where
-- 到着駅ごとに列車を仕分ける
    trains = accumArray (flip (:)) [] (1,n) [(b, t) | t@(_:_:_:_:_:b:_) <- ldkcabs]
-- 終電の出る時刻を順に調べる
    ans = loop IM.empty $ PQ.singleton $ PQ.Entry (- maxBound) n
-- キューの先頭から順に調べる
    loop im q
      | PQ.null q = im  -- キューが空なら終わり
      | IM.member b im = loop im q1 -- bは確定済みならスルー
      | otherwise = loop (IM.insert b (- nt) im) (PQ.union q1 q2)
      where
        Just (PQ.Entry nt b, q1) = PQ.uncons q
        q2 = PQ.fromList
             [ PQ.Entry (negate $ l + j * d) a -- aからbに向かう-ntに間に合う出発時刻
             | l:d:k:c:a:_ <- trains ! b
             , let j = min (pred k) $ div (- nt - l - c) d, 0 <= j]
```

# F - Black Jack

[問題 ABC342F](https://atcoder.jp/contests/abc342/tasks/abc342_f)

シグネチャを決める。

```haskell
import Data.Array

abc342f :: Int    -- N
        -> Int    -- L
        -> Int    -- D
        -> Double -- 答え
abc342f n l d = ...
```

（方針は合っていたが混乱して完成させられなかったので、フレンズさんのヒントを見て話を整理することでようやく正解に至った。）

## 考える

いくつか値を定義しておく。

```haskell
    ld1 = pred l + d        -- L + D - 1
    nd1 = pred n + d        -- N + D - 1
    nd2 = succ n + succ d   -- N + 1 + D + 1
    dd = fromIntegral d     -- Double型のD
```

### ディーラーの手がyになる確率 p[y]

ディーラーの手順は、プレイヤーの結果$x$がいくつかには影響されない。なので独立して計算できる。

ディーラーは手が$L$未満なら引き続きサイコロを投げる。
$p[y_1] = p_1$が決まったとき、$p[y_1+1]$から$p[y_1+D]$までに$p_1/D$を足し込む、という操作を$p[0] = 1$から始めて$p[L-1]$までを配るDPで終えたときの、$p[L]$から$p[L+D-1]$が求める確率である。これらの和は1になる。

しかしこの配るDPでは足し算の回数が$O(LD)$となり多すぎる。
集めるDPで考えると、$p[y]$はその手前$1～D$からそれぞれ$1/D$の確率で到達するので、手前$D$要素の和を$D$で割った値になる。「手前$D$要素の総和」を今作った値と、$D+1$個前の値とを使って尺取法のように更新して進めることで、$p[y]$を全て求めることができる。なお、$L$に達した後は足すものはなくなる。

```haskell
    p = listArray (0, ld1) $ 1 : pLoop 1 0
    pLoop y acc = acc1 / dd : pLoop (succ y) acc1
      where
        incr = if y <= l then p ! pred y else 0       -- Lまでは直前の値を足し込む
        decl = if d <  y then p ! (pred y - d) else 0 -- p[0]以前は0を引く
        acc1 = acc + incr - decl
```

ゲームにおいてはディーラーの手が$L$未満で止まることないので、$p[0]$から$p[L-1]$は0で上書きする。
```math
p'[x] = \left \{
\begin{array}{ll}
0 & (0 \leq x < L) \\
p[x] & (L \leq x < L + D)
\end{array}
\right .
```

### プレイヤーの手がxのときの勝率q[x]

ゲームのルールは、$x \leq y \leq N$のときプレイヤーの負け、それ以外のとき勝ち、である。
プレイヤーの出した手 $x$ に対してディーラーは $y$ を $x$ から $N$ の間に入れようとするので、プレイヤーはなるべくそれを狭くすることが勝つ確率を上げることになる。ただし攻めすぎて$N$を越えたらバーストで即負けする。

$0 \leq x \leq N$の範囲でプレイヤーの勝つ確率を考える。
バーストのときは負け。$q[N < x] = 0$
そうでないとき、$x \leq y \leq N$ とならない確率。つまり $y < x$または$N < y$となる確率。$q[0 \leq x \leq N] = \sum_{y=0}^x p'[y] + \sum_{y=N+1}^{N+D} p'[y]$
第2項は定数。ディーラーがバーストする確率。これをベーシックインカム`bi`と呼ぼう。

```haskell
    bi = sum [p ! y | y <- [succ n .. ld1]]
```

第1項は累積和で求められる。下の方は0行進であるが、気にせず計算しよう。

```haskell
    p' x = if l <= x && x <= ld1 then p ! x else 0
    q = listArray (0, nd1) $
        scanl (+) bi [p' x | x <- [0..pred n]] ++  -- 累積和の初期値をBIにした
        repeat 0                                   -- バースト
```

### 戦略

勝率を最大にする最適な戦略をとったとき、手$x$から勝つ確率を$a[x]$とする。
勝率を最大にするには、現在の手が $x$ であるとき、

- ここで止める。勝率 $q[x]$
- さらに続ける。勝率は、$a[x+1] \sim a[x+D]$の平均値

の大きい方を選ぶ。$a[x] = \max(q[x], \frac{1}{D} \sum_{i=1}^D a[x+i])$
また$D$要素の総和がでてきて、今度は後ろから尺取法をして$a[\cdot]$を求めていく。

```haskell
    a = array (0, nd2) $
        [(x, 0) | x <- [succ n..nd2]] ++ -- バースト
        aLoop n 0
    aLoop (-1) _ = []
    aLoop x acc = (x, max (q ! x) (acc1 / dd)) : aLoop (pred x) acc1
      where
        acc1 = acc + a ! succ x - a ! (x + succ d)
```

（直前で作ったばかりの `a ! succ x` をわざわざ配列から読み出しているのがぎこちない。`pLoop`は場合分けがあったがここにはないので、もっと普通に書けるはず。）

### 答え

$a[0]$が問題の答えになる。

```haskell
abc342f n l d = a ! 0
  where
    ...
```

## 公式解説から

配るDPでやるには、$p[y]$の値を$p[y+1]$～$p[y+D]$に足し込むために、$y+1$に$+p[y]$ $y+D+1$に$-p[y]$という増分を与えて累積和をとる、いわゆるいもす法のやり方をすればいい、とあった。なるほど。

# G - Retroactive Range Chmax

[問題 ABC342G](https://atcoder.jp/contests/abc342/tasks/abc342_g)

シグネチャを決める。
セグメント木の操作が必要なので、全体をIOアクションにする。

```haskell
abc342g :: Int     -- N
        -> [Int]   -- Ai
        -> Int     -- Q
        -> [[Int]] -- query_i
        -> IO ()
abc342g n as q qus =
  do
    ...
  where
    ...
```

## 考える

- 配列の初期値を設定する
- 配列のとある区間に同様の操作を施す
- 配列の一つの要素を読み出す

これは、いわゆる区間更新セグメント木の機能の一部である。
そこで、上の3つの機能だけを持ったダイジェスト版を作ることにする。

書き換え可能なメモリとして、区間更新セグメント木の区間更新側の情報を持つ$2N-1$要素の配列が必要である。`Data.Array.IO` で実装しよう。

```haskell
  do
    arr <- newArray (0, w + pred w) *** :: IO (IOArray Int ***) -- 更新情報の型未定
    ...
  where
    w = head $ dropWhile (n >) $ iterate (2 *) 1 -- セグメント木配列の要素数
```

$A_i$ の値を保持する配列は更新しないので `Data.Array` でいい。
タイプ2のために、クエリも配列に入れておく。

```haskell
    a = listArray (1,n) as
    qu = listArray (1,q) qus
```

## 区間更新の情報

タイプ1のクエリは、値をクリップする下限を指定する。同じ範囲にいくつもの値が重なる可能性があるが、タイプ3のクエリに応じて使われる値はその最大値だけである。
タイプ2のクエリで取り消されることもあるので、全てを持っておく。

ここでは、指定された値の降順のリストを保持することにする。

- タイプ1のクエリでは、`insertBy (flip compare)` で降順を保って挿入する
- タイプ2のクエリでは、`delete` で一つ削除する
- タイプ3のクエリでは、最大値は先頭にあるのでそれとの `max` をとる

とすればよい。

他の選択肢として、値の多重集合をつかう生真面目なやり方や、整列しない値リストと最大値の対を使う方法なども考えられる。

多重集合を使うならば、
- タイプ1では、要素を一つ加える
- タイプ2では、要素を一つ除く
- タイプ3では、実装に使う`Data.IntMap.findMax`で最大値を取り出す

整列しないリストならば、
- タイプ1では、要素をconsする。最大値は新しい値との`max`で更新
- タイプ2では、`delete` で削除し、リスト全体を対象に `maximum` を取り直す
- タイプ3では、求めておいた最大値を使う

ここで、`maximum`の取り直しは、遅延評価されるので、タイプ2が連続したような場合では、使われない値は計算されずに捨てられる。

全てやってみたが、降順リストか最速だった。

### クエリ対応

クエリのタイプごとに分岐する。

```haskell
    forM_ qus (\qi ->
      case qi of
        1:l:r:x:_ -> {- タイプ1の処理 -}
        2:i:_     -> {- タイプ2の処理 -}
        3:j:_     -> {- タイプ3の処理 -}
      )
```

### 読み出し

タイプ3のクエリに対しては、$A_i$ を初期値として、セグメント木を葉から根に辿り、先頭の値でクリップすることを繰り返せばよい。

起動時には、初期値を配列から読み出し、処理の本体を呼び出す。結果を `print` する。

```haskell
        3:j:_     -> type3 w arr j (a ! j) >>= print
```

添え字が0始まりにずれているので、第$j$要素のセグメント木配列での添え字は$w-1+j-1$から始まり、1引いて2で割ると親に行ける。

```haskell
type3 :: Int -> IOArray Int [Int] -> Int -> Int -> IO Int
type3 w arr j v0 = loop (pred w + pred j) v0
  where
    loop (-1) v = return v
    loop i v = do
      xs <- readArray arr i
      let v1 = if null xs then v else max v (head xs)
      loop (div (pred i) 2) v1
```

### 区間更新

区間 $[a,b)$ の更新値に対する操作を関数でわたして、区間更新配列を修正する処理。
現在のノードが指定範囲から外れているとき何もしない、完全に含まれているとき操作を適用、部分的に交わっているとき、左右の両方の部分木に再帰する。
現在のノードが担当する範囲は再帰呼び出しの間追跡する。

```haskell
-- [a,b)にfを施す
rangeUpdate :: Int -> IOArray Int a -> Int -> Int -> (a -> a) -> IO ()
rangeUpdate w arr a b f = loop 0 w 0
  where
    loop p w i
      | q <= a || b <= p = return ()
      | a <= p && q <= b = readArray arr i >>= writeArray arr i . f
      | otherwise = loop p w2 (i + succ i) >> loop (p + w2) w2 (2 * succ i)
      where
        q = p + w
        w2 = div w 2
```

### タイプ1

タイプ1は、$[l,r)$の区間に$x$を降順を保って挿入する。
上の `rangeUpdate` を用いて直接定義できる。

```haskell
        1:l:r:x:_ -> rangeUpdate w arr (pred l) r (insertBy (flip compare) x)
```

### タイプ2

タイプ2は$i$番目のクエリの内容の$l,r,x$について、それを除去する。
こちらも `rangeUpdate` を用いて直接定義できる。

```haskell
        2:i:_     -> let _:l:r:x:_ = qu ! i in rangeUpdate w arr (pred l) r (delete x)
```

結果：
[1821ms, 328MB](https://atcoder.jp/contests/abc342/submissions/50947913)
参考：多重集合を使った場合
[3218ms, 449MB](https://atcoder.jp/contests/abc342/submissions/50707522)
単純なリストと最大値の対を使った場合
[4742ms, 506MB](https://atcoder.jp/contests/abc342/submissions/50707147)

## 公式他の解説について

汎用性のあるデータ構造にこの問題を押し込もうとすると、「双対セグメント木」というものが必要になって、いろいろ大変なようだ。