ABC335をHaskellで

# A - 202\
3\


[問題 ABC335A](https://atcoder.jp/contests/abc335/tasks/abc335_a)

```haskell
main = getLine >>= putStr . init >> putStrLn "4"
```

# B - Tetrahedral Number

[問題 ABC335B](https://atcoder.jp/contests/abc335/tasks/abc335_b)

シグネチャを決める。3項組はタプルで表すことにする。

```haskell
abc335b :: Int  -- N
        -> [(Int,Int,Int)] -- 答え (x,y,z)
```

いろいろな型と次元の扱える `Data.Array` 配列の添え字を一次元の0以上の整数に写像する機構 `Data.Ix` を使うと手抜きできる。

## 結果

```haskell
import Data.Ix

abc335b n = filter (\(x,y,z) -> x + y + z <= n) $ range ((0,0,0),(n,n,n))
```

## 素朴な解

```haskell
abc335b n =
  [ (x,y,z)
  | x <- [0..n]
  , y <- [0..n - x]
  , z <- [0..n - x - y]
  ]
```

# C - Loong Tracking

[問題 ABC335C](https://atcoder.jp/contests/abc335/tasks/abc335_c)

シグネチャを決める。
クエリは $1 \\; C$ を `Left Char` $2 \\; p$ を `Right Int` とする `Either Char Int` で表す。（専用のデータ型を作るのをサボって）

```haskell
type Query = Either Char Int
type PT = (Int,Int)

abc335c :: Int     -- N
        -> Int     -- Q
        -> [Query] -- query_i
        -> [PT]    -- 答え
```

## リングバッファな解法

パーツ$i$の番号を添え字にした配列だと、龍が進むたびに$N$要素を転送することになって辛い。
また、末尾は順に不要になっていくので、要素数$N$のリングバッファを使って、
クエリ2については、現在の頭の位置から数える、とやると考えた。

[結果](https://atcoder.jp/contests/abc335/submissions/49135389)は泣きたくなるくらいに命令型のプログラムになった。
書き換えできてランダムアクセス可能な配列が中心なので仕方ないことだが、あまりにアレなので違う方法を考えてみる。

## 頭の位置を全て記録する解法

クエリを前から順に計算しようとすると、クエリ2でランダムアクセスな配列が必要になる。
クエリ列全体が既に確定していることを活用して少しズルをする。
時刻0における頭の位置は(1,0)で、それ以降、クエリ1の度に頭は移動する。
この頭の位置の推移を、クエリ1全てに対して記録しておく。
クエリ2は、現在より$p-1$だけ以前の頭の位置を取り出す。
それは、このクエリ2以前に出現したクエリ1が$i$個であったとすると、記録の前から$i - p + 1$番目になる。
なお、このインデックスがはみ出すときは、初期状態の胴体の位置を調べているので、固定的に計算できる。

```haskell
abc335c :: Int -> Int -> [Query] -> [PT]
abc335c n q qus = loop 0 qus
  where
    q1s = lefts qus
    pts = listArray (0, length q1s) $ scanl move (1,0) q1s
    loop _ [] = []
    loop i (Left _ :qus) = loop (succ i) qus  -- クエリ1の個数が増える
    loop i (Right p:qus) = xy : loop i qus
      where
        j = i - pred p
        xy | j >= 0    = pts ! j
           | otherwise = (1 - j, 0)

move :: (Int,Int) -> Char -> (Int,Int)
move (x,y) 'R' = (succ x, y)
move (x,y) 'L' = (pred x, y)
move (x,y) 'U' = (x, succ y)
move (x,y) 'D' = (x, pred y)
```

これは[ユーザ解説 by cirno3153](https://atcoder.jp/contests/abc335/editorial/9053)と同じで、動的配列を使う代わりにクエリ1を全て済ませてからクエリ2を処理しているだけといえる。

フレンズさんの[公式解説](https://atcoder.jp/contests/abc335/editorial/9026)の末尾にさらっと書かれている

> 配列をリバースし「先頭の要素を取り除く」処理を実際には行わない

も同じ事を言っている。

# D - Loong and Takahashi

[問題 ABC335D](https://atcoder.jp/contests/abc335/tasks/abc335_d)

シグネチャを決める。
0を`T`、それ以外の整数を龍とする$N \times N$配列で結果を返す。

```haskell
abc335d :: Int -- N
        -> Array (Int,Int) Int -- 答え
```

例1のように、左上から始めて、グルグルとトグロを巻いていけばよい。
$N$は奇数なので、一周巻くと$N-2$行$N-2$列が残る。

位置 $(k,k)$ からサイズ $m+1$ のトグロを $s$ から始めて一周巻く、を再帰的に繰り返す。

## 結果

```haskell
abc335d n = array ((1,1),(n,n)) $ recur 1 (pred n) 1
  where
    recur k 0 _ = [((k, k), 0)]
    recur k m s =
      recur (succ k) (m - 2) (s + 4 * m) ++
      [((k     + i, k        ), s         + i) | i <- [0 .. m]] ++
      [((k + m    , k     + i), s +     m + i) | i <- [1 .. m]] ++
      [((k + m - i, k + m    ), s + 2 * m + i) | i <- [1 .. m]] ++
      [((k        , k + m - i), s + 3 * m + i) | i <- [1 .. m - 1]]
```

# E - Non-Decreasing Colorful Path

[問題 ABC335E](https://atcoder.jp/contests/abc335/tasks/abc335_e)

シグネチャを決める。

```haskell
abc335e :: Int         -- N
        -> Int         -- M
        -> [Int]       -- Ai
        -> [(Int,Int)] -- Ui, Vi
        -> Int         -- 答え
```

整数が等しい頂点だけを通って行き来できる頂点群を一つの頂点と見なしてグラフを組み直す。
これはUnion-Findで分割の代表元を使えばできる。
組み直したグラフで、頂点1のスコアを0、それ以外は、整数が真に小さい隣接頂点のスコアの最大値+1を値とする、集めるDPにより、頂点$N$のスコアを求める。

## 結果

[全体](https://atcoder.jp/contests/abc335/submissions/49137791)

```haskell
import Control.Monad
import Data.Array
import Control.Monad.ST

abc335e :: Int -> Int -> [Int] -> [(Int,Int)] -> Int
abc335e n _m as uvs = max 0 $ s ! (roots ! n)
  where
    a = listArray (1,n) as -- Ai
    roots = runST $ do -- 頂点uの属する分割の代表元を持つ配列
      uf <- newUF (succ n)
      forM_ uvs (\(u,v) -> when (a ! u == a ! v) $ uniteUF uf u v)
      listArray (1,n) <$> forM [1..n] (getRoot uf)
    g = accumArray (flip (:)) [] (1,n) -- 代表元だけ、Aの大きい方から小さい方への辺だけのグラフ
        [ if au > av then (u1, v1) else (v1, u1)
        | (u, v) <- uvs
        , let u1 = roots ! u, let v1 = roots ! v, u1 /= v1
        , let au = a ! u, let av = a ! v
        ]
    s = listArray (1,n) $ map sf [1..n]
    v1 = roots ! 1
    sf v
      | v == v1 = 1
      | null us = minBound
      | otherwise = succ $ maximum $ map (s !) us
      where
        us = g ! v

-- UnionFind (実装は省略)

-- 作成
newUF :: Int -> ST s UnionFind
-- 代表元
getRoot :: UnionFind -> Int -> ST s UnionFind
-- 統合
uniteUF :: UnionFind -> Int -> Int -> ST s ()
```

# F - Hop Sugoroku

[問題 ABC335F](https://atcoder.jp/contests/abc335/tasks/abc335_f)

シグネチャを決める。

```haskell
abc335f :: Int   -- N
        -> [Int] -- Ai
        -> Int   -- 答え
abc335f n as = ...
```

（わからなくて解説を見た。）

## 公式解説の復唱

まず、黒く塗るマスとは結局、スゴロクで止まったマスのこと。ゴールしたときに振り返って、マスの塗り方の場合の数とはつまり、ゴールに到達する経路の種類数と同じ。（まずこの理屈に気がつかなかった…）

そうと決まれば、マスに場合の数を書き込んで、配るDPなり集めるDPなりをしようと考える。マス$i$への到達経路の個数を$C[i]$とする。$C[1]=1$で、いつ終わってもいいので$\sum_i C[i]$が答え。
ただこのスゴロクの進め方が変わっていて、マス$i$ごとに飛べる次のマスの間隔$A_i$が異なり、しかもその整数倍ならいくらでも遠くまで進むことが許される。

### 集めるDP (失敗)

集めるDPで考えると、$i$より前の全てのマス$1 \leq j < i$について、距離の差が$A_j$の倍数なら移動できるので $C[i] = \sum_{j \in S_i} C[j]$ ここで $S_i = \\{j \\ |\\  1 \leq j < i, (j - i) \bmod A_j = 0 \\}$ となる。$S_i$の要素を調べるために手前全てのマスを確認する必要があり、$O(N^2)$になるので無理。

```haskell
abc335f n as = sum $ elems c
  where
    a = listArray (1,n) as
    c = listArray (1,n) $ map cf [1..n]
    cf 1 = 1
    cf i = sum [c ! j | j <- [1 .. pred i], mod (i - j) (a ! j) == 0]
```

### 案1 : 普通に配るDP

配るDPで考えると、動作はエラトステネスの篩じみて、$C[i]$をその先のマスに足し込む。
$C[j \in T_i] \leftarrow C[j] + C[i]$ ここで $T_i = \\{ j \\ | \\ j = i + xA_i, 0 < x \leq \frac{N-i}{A_i} \\}$ $x$ の制約は $i < j \leq N$ のこと。

```
C[1～N] = 0 -- カウント配列 添え字はマスの番号
C[1] = 1

for i in [1 .. N-1]                  -- 順にC[i]を配っていく
   for j in [i + A[i], i + 2A[i] .. N]
      C[j] += C[i]

return (sum C[])
```

普通のサイコロなら目の上限で配る先の要素数は抑えられるが、この問題設定だと$|T_i|$が$N$まで増える恐れがある。話が前後するけどABC342Fの$D$が大きい場合と似ている。あちらは配る先が連続区間なので、いもす法を使うという策がとれた。こちらは配る先が飛び飛びなのでそれもできない。

### 案2 : ローラースタンプに預ける、配るDP

ここで（解説にある）解決策は、配る周期$p$ごとに配る値を$p$マス抱えるローラースタンプを用意し、直接$C[i]$に値を配る代わりに、狙ったマスに投下するようなスタンプのマスに値を預けるというもの。注目するマスを$i$に進めるときは、ローラースタンプを転がして、そのマスに押すべき全ての印を押すことで、配る操作が完了する。その値は$A_i$周期のローラースタンプのマスに足し込まれる。

```
C[1～N] = 0
C[1] = 1

S[1～N] = [] -- ローラースタンプの束 添え字は周期periodの長さ
for p in [1 .. N]
   S[p] = (SS[0～p-1] = 0) -- 周期pのローラースタンプ 添え字はマスの番号 mod p

for i in [1..N-1]  -- 順にC[i]を処理していく

   for p in [1 .. N]          -- スタンプに預けられた値をC[i]へ「配る」
      C[i] += S[p][i mod p]   -- (*)

   S[A_i][i mod A_i] += C[i]  -- 以降で配るべき値をスタンプに預ける

return (sum C[])
```

上手い考えに見えるが、束のスタンプ$N$本を全てのマスに毎回押す `(*)` の処理が $O(N^2)$ になる。

### 折衷案

もともと、案1でも、$A_i$の値がそれなりに大きいときは、配る相手も限られるため機能する。
$A_i$の値が小さくて、直接配ると配る先が多くなりすぎる場合への対処が案2なので、両方を同時に使う。
つまり、境界 $B$ を定め、$A_i \leq B$ のものは案2のスタンプ送りにし、$B < A_i$ のものは直接配る。

スタンプの本数$B$は`(*)`の処理の回数で、スタンプを使わず直接配るマスの個数は最大$N/B$なので、全体の計算量は $O(N \cdot \max(B, N/B))$ となる。$B= \sqrt N$とするとこれを最小化できて$O(N \sqrt N)$ となるので、平方分割するということ。

```
B = floor (sqrt N) -- 境界

C[1～N] = 0
C[1] = 1

S[1～B] = []               -- 要素数修正
for p in [1 .. B]
   S[p] = (SS[0～p-1] = 0)

for i in [1..N-1]

   for p in [1 .. N]        -- (スタンプに預けられた値をC[i]へ配る)
      C[i] += S[p][i mod p]

   if A_i ≦ B                   -- 境界で対応を切り替える
      S[A_i][i mod A_i] += C[i]  -- (以降で配るべき値をスタンプに預ける, 案2モード)
   else
      for j in [i + A[i], i + 2A[i] .. N] -- (直接配る, 案1モード)
         C[j] += C[i]

return (sum C[])
```

[STArrayで書いた版, 761ms, 29MB](https://atcoder.jp/contests/abc335/submissions/49149950)
スタンプ束を三角配列で作るのをサボってメモリを無駄に使っている。

# G - Discrete Logarithm Problems

[問題 ABC335G](https://atcoder.jp/contests/abc335/tasks/abc335_g)

シグネチャを決める。

```haskell
abc335g :: Int   -- N
        -> Int   -- P
        -> [Int] -- Ai
        -> Int   -- 答え
abc335g n p as = ...
  where
    ...
```

## 考える

[アライグマ「G問題は数学なのだ。頑張るのだ！」](https://twitter.com/kyopro_friends/status/1743635494569705630)と言われてもなんもわからんので公式解説をなぞる。

まず、$P-1$を素因数分解して、素因数の一覧を求める。

```haskell
    ps = map head $ group $ primeFactors (pred p)

primeFactors :: Int -> [Int] -- 素因数分解して昇順のリストで返す (primeFactors 12 => [2,2,3])
```

$P$に関する位数を求める関数`rank`を定義する。

解説のitemizeを命令型の擬似コードにするとこう：
```
function rank(x) -- xの位数を求める
   M = P - 1
   for pk in ps -- P-1の素因数リスト
     while (M % pk == 0 ∧ x^(M / pk) == 1 mod p)
        M = M / pk
   return M
```

Haskellだとこう：
```haskell
    rank :: Int -> Int
    rank x = foldl step (pred p) ps
      where
        xl = fromIntegral x
        step m pk
          | r == 0, cond = step q pk
          | otherwise    = m
          where
            (q,r) = divMod m pk
            cond = pow xl q == 1

    pow :: Integer -> Int -> Integer
    pow a b = powerish mul 1 a b
    mul :: Integer -> Integer -> Integer
    mul x y = mod (x * y) (fromIntegral p)

powerish :: Integral p => (b -> b -> b) -> b -> b -> p -> b
powerish op i a b = ... -- i `op` a `op` a `op` ... `op` a  ようは i * a^b を求める
```

普段の問題なら64ビットIntの計算で間に合うが、この問題は$P \leq 10^{13}$が大きくてオーバーフローするので、$x^{M/p_k} = 1 \bmod P$ の判定を行う `cond` 回りだけ `Integer` で計算するようにしている。

そして$A_i$の位数$B_i$を全て求め、$B_j | B_i$ を満たす組み合わせの個数を求める。
そのためにまず、位数の値ごとにその個数を数えて貯める。
$P$は大きいので配列でなく`IntMap`を使う。

```haskell
    bee = IM.fromListWith (+) [(rank a, 1) | a <- as]
```

$B_i$が$B_j$を割りきるような組み合わせについて、個数の積を足し合わせたら答え。

```haskell
abc335g _n p as = sum
    [ c * d
    | bics@((bi, c):_) <- tails $ IM.assocs bee
    , (bj, d) <- bics
    , mod bj bi == 0
    ]
```

### 位数の計算を速くする

[ユーザ解説 by MMNMM](https://atcoder.jp/contests/abc335/editorial/9042)のitemizeを擬似コード化する。
```
function rank(x) -- xの位数を求める
   M = P - 1
   for pk in ps -- P-1の素因数リスト
     while (M % pk == 0)
        M = M / pk
     t = x ^ M
     while (t ≠ 1 mod p)
        M = M * pk
        t = t ^ pk
   return M
```

公式解説のものと見比べてみると、いずれにせよ $M$ は初期値 $P-1$ の素因数をいくつか無くした値$\prod p_k^{a \leq e_k}$ をさまようことは同じであるとわかる。
公式のやり方は、$M$ の素因数に $p_k$ がなくなったら強制終了し、そうでない範囲で、$x ^ M = 1 \bmod P$ を満たす最小の $M$ を、$p_k$ で割ってだんだん小さくしていくことで発見しようとしている。
ユーザ解説の方法は、最初に $M$ を $p_k^{e_k}$ で（とは書いてないけどそういうこと）割りきってしまい、$x ^ M = 1 \bmod P$ になるまで、$M$ に $p_k$ を掛けて最小の数を発見しようとしている。
つまりこの条件は単調なのだろう。

割り算は四則演算の中では他より重いので「$M$ を $p_k$ で割れるだけ割り」は優しくない。$P-1$ を素因数分解したときに $p_k^{e_k}$ も求めておけば、割り算は一度で済む。

```haskell
-- 素因数p_kと、あるだけ掛けたp_k^e_kの対のリスト
    pes = map (\xs -> (head xs, product xs)) $ group $ primeFactors (pred p)

-- 位数を求める関数
    rank :: Int -> Int
    rank x = foldl step (pred p) pes
      where
        xl = fromIntegral x
        step m (pk, pkek) = loop m0 (pow xl m0)
          where
            m0 = div m pkek
            loop m 1 = m                          -- t == 1 になるまで繰り返し
            loop m t = loop (m * pk) (pow t pk)   -- Mにpkを掛け、tをpk乗する
```

最初に書いた未整理版は[2392ms](https://atcoder.jp/contests/abc335/submissions/49155184)と確かに速くなったのに、これを書くために整理した版は[4472ms](https://atcoder.jp/contests/abc335/submissions/51008127)でむしろ遅くなってしまった。何が違うのか訳がわからない。