ABC262 A～E をHaskellで

# A - World Cup

[問題 ABC262A](https://atcoder.jp/contests/abc262/tasks/abc262_a)

まず、関数の型を次のように決める。

```haskell
abc262a :: Int -- Y
        -> Int -- 答え
```

リスト内包表記を使えば、$Y$ 年から順に+1年ずつ、年を4で割った余りが2であるような年だけを順に取り出すことは簡単に書ける。（$Y$ 年から $Y+3$ 年の間に必ず答えが1つだけあるので生成器の終値を設定してもよいが、面倒なので省略する。）

```haskell
[k | k <- [y ..], mod k 4 == 2]
```

このリストの先頭が求める答えである。

## 結果

```haskell
abc262a :: Int -> Int
abc262a y = head [k | k <- [y ..], mod k 4 == 2]
```

# B - Triangle (Easier)

[問題 ABC262B](https://atcoder.jp/contests/abc262/tasks/abc262_b)

関数のシグネチャを定める。

```haskell
abc262b :: Int -> Int    -- N, M
        -> [(Int,Int)]   -- (UiとViのペア)のリスト uvs
        -> Int           -- 答え
```

全体としては、条件を満たす組み合わせを、$a, b, c$ を総当たりして数え上げることで答えを得られる。

```haskell
length [ ()
       | a <- [1 .. n-2]
       , b <- {- aと結ばれている頂点 -}
       , c <- {- bと結ばれている頂点 -}
       , {- aとcが結ばれている -}
       ]
```

これが動くようにするには、ノード番号をキーとして、接続されている（より大きい）ノード番号の一覧を取得できるようにする必要がある。毎回 `uvs` をスキャンして
```haskell
[v | (u,v) <- uvs, a == u] -- aと結ばれている頂点
```
などとするのは計算機に酷なので、そのような配列を`Data.Array.accumArray` で先に構築しておくとよい。

```haskell
edges = accumArray (flip (:)) [] (1,n) uvs
```

## 結果

```haskell
import Data.Array

abc262b :: Int -> Int -> [(Int,Int)] -> Int
abc262b n m uvs = length
    [ ()
    | a <- [1 .. n-2]
    , b <- edges ! a
    , c <- edges ! b
    , elem c (edges ! a)
    ]
  where
    edges = accumArray (flip (:)) [] (1,n) uvs
```

<!--
## 続き

リスト内包表記でやっていることをよく見ると、それぞれの`a`に対して、`edges ! a` と `edges ! b` の共通する要素の個数を求めている。

```haskell
import Data.List

abc262b n m uvs = sum
    [ length $ intersect (edges ! a) (edges ! b)
    | a <- [1 .. n-2]
    , b <- edges ! a
    ]
  where ...
```
-->

# C - Min Max Pair

[問題 ABC262C](https://atcoder.jp/contests/abc262/tasks/abc262_c)

シグネチャを決める。
```haskell
abc262c ::  Int   -- N
        -> [Int]  -- Ai
        ->  Int   -- 答え
```

数列の異なる位置 $i,j$ から値を取り出して（$a_i, a_j$）、いずれかが $i$ と等しく、いずれかが $j$ と等しくなる場合とは、（$i < j$ より $i \neq j$ であることを考慮して）

- $a_i = i, a_j = j$ : 要素が位置と同じ値なものの二つ組
- $a_i = j, a_j = i$ : 要素は位置と異なる値で、ちょうど互いを指す二つ組

のいずれかである。

前者は、位置と同じ値な要素の個数を $x$ としたとき、そこから2個選ぶやり方の個数なので ${}\_xC\_2$ となる。

```haskell
let
  x = length [() | (i,ai) <- zip [1..] as, i == ai]
in
  div (x * pred x) 2
```

後者は、位置と異なる値なそれぞれの要素 $a_i = j$ に関して、$a_j = i$ と戻ってくるかどうかは $j$ の1箇所を確認すればわかる。ここで条件を $a\_i \neq i$ だけにすると倍数えてしまうので、 $a\_i > i$ として組を一度だけ数えるようにする。
それと、$a_j$ を取り出すのに `as !! pred j` とリストのままするのは遅いので、配列に移しておくべきである。

```haskell
let
  aa = listArray (1,n) as
in
  length [() | (i,ai) <- zip [1..] as, i < ai, aa ! ai == i]
```

## 結果

```haskell
import Data.Array

abc262c :: Int -> [Int] -> Int
abc262c n as = div (x * pred x) 2 + y
  where
    x = length [() | (i,ai) <- zip [1..] as, i == ai]
    aa = listArray (1,n) as
    y = length [() | (i,ai) <- zip [1..] as, i < ai, aa ! ai == i]
```
<!--
## 難読化

`x`と`y`を一度の走査で数えるには、

- $i = a_i$ なら`x`を1増やす
- そうでなく、$i < a_i$ かつ $a_{a_i} = i$なら`y`を1増やす
- そうでもないならどちらも増やさない

と場合分けして

```haskell
(xs,ys) = unzip $ zipWith check [1..n] as
check i ai
  | i == ai                 = (1,0)
  | i <  ai && aa ! ai == i = (0,1)
  | otherwise               = (0,0)
(x,y) = (sum xs, sum ys)
```

とできる。

```haskell
(_,x,y) = foldl step (1,0,0) as
step (i,x,y) ai
  | i == ai                 = (x+1,y)
  | i <  ai && aa ! ai == i = (x,y+1)
  | otherwise               = (x,y)
```
-->

# D - I Hate Non-integer Number

[問題 ABC262D](https://atcoder.jp/contests/abc262/tasks/abc262_d)

シグネチャを決める。

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

モジュロ演算の足し算を作っておく。
```haskell
modBase = 998244353
add x y = mod (x + y) modBase
```

$N \leq 100$ と問題の規模は小さいようにみえるが、全ての場合の数は問題文にあるように $2^N-1$ もあるので、総当たりで生成するやり方では途方もなく時間がかかる。なので動的プログラミングを用いる。

全てを一度に解くのは無理そう。
部分問題として、数列から整数をちょうど $k$ 個選択したとき、平均値が整数である場合の数を数えることを考える。これができれば、$k = 1,\dots,N$ について求めた結果の総和が答えとなる。

「整数 $k$ 個の平均値が整数である」とは「整数の和が $k$ の倍数である」ことに等しい。
数列 $Ai$ の前から $x \  (0 \leq x \leq N)$ 個だけを考え、その中から $y \  (0 \leq y \leq k)$ 個を選択した和を $k$ で割った余りが $z \  (0 \leq z < k)$ になるような選び方の個数を追跡する。
配列の変数名をカウントの $cnt$ とする。$y$ と $z$ の2次元配列を $x$ ごとに作るので、$cnt_x[0\leq y \leq k][0 \leq z < k]$ のようになる。

$cnt_x$ の $y=0$ の行は、和は0にしかならないので $cnt_x[0][0]=1, cnt_x[0][z\neq 0] = 0$ で固定となる。この行は省略できる。$cnt_0$ については $y > 0$ の全ての行も0である。

$cnt_{x-1}$ と $a_x$ から $cnt_x$ が計算できる。
- $a_x$ を使わない場合について、$cnt_{x-1}$ の値が全てそのまま $cnt_x$ に引き継がれる。
- $a_x$ を使う場合について、使う数がひとつ増えるので $y$ を1増やしたところに、和は $a_x$ だけ増えるのでそれだけずらした位置に、場合の数をさらに足しこむことになる。

つまり：

$cnt_x[y][(z + a_x) \bmod k] = cnt_{x-1}[y][(z + a_x) \bmod k] + cnt_{x-1}[y-1][z]$
$(cnt_{x-1}[0][z] = [1,0,\dots,0]とする)$

$cnt_x$ をリストのリストで表し、`zipWith`の組み合わせで上の漸化式を実現するとかっこいいが、リストセルのメモリ割り当てに時間がとられてTLEしてしまうため、素直に `Data.Array.accum` を用いる。

```haskell
cntn = foldl step cnt0 as
cnt0 = accumArray (+) 0 ((1,0),(k, pred k)) []
step cnt ax = accum add cnt $
  ((1, mod ax k), 1) :
  [((y, mod (z + ax) k), cnt ! (pred y, z)) | y <- [2..k], z <- [0..pred k]]
```

数列の全ての要素を処理した後の、$cnt_N[k][0]$ が求める値である。

```haskell
f k = cntn ! (k,0)
  where ...
```

特別な場合として、$S=Σ A_i$ とおいて
$f(1) = N$
$f(2) = {}_eC_2 + {}_oC_2$ （$e$ は $A_i$ の偶数の個数、$o = N - e$ は奇数の個数）
$f(N) = 1$ （$S$ が $k$ で割り切れるとき）
$f(N) = 0$ （割り切れないとき）
$f(N-1) = $ $S-A_i$ のうち $N-1$ で割り切れるものの個数
はどれも $O(N)$ で、上述の手順を走らせるよりも速く結果を求められるが、そこを凝る必要もないだろう。

## 結果

```haskell
import Data.Array

modBase = 998244353
r x = mod x modBase
add x y = r (x + y)

abc262d :: Int -> [Int] -> Int
abc262d n as = r $ sum $ map f [1..n]
  where
    f k = cntn ! (k,0)
      where
        cntn = foldl step cnt0 as
        cnt0 = accumArray (+) 0 ((1,0),(k, pred k)) []
        step cnt ax = accum add cnt $
          ((1, mod ax k), 1) :
          [((y, mod (z + ax) k), cnt ! (pred y, z)) | y <- [2..k], z <- [0..pred k]]
```
# E - Red and Blue Graph

[問題 ABC262E](https://atcoder.jp/contests/abc262/tasks/abc262_e)

シグネチャを決める。

```haskell
abc262e :: Int -> Int -> Int   -- N, M, K
        -> [ (Int,Int) ]       -- Ui, Vi のベアのリスト uvs
        -> Int                 -- 答え
```
モジュラな演算なども定義しておく。
```haskell
p = 998244353
r x = mod x p
mul x y = r (x * y)
```

ノードは最初全て青く塗られているとし、そのうちいくつか選んで赤く塗り替えていくことを考える。下の図で示すとおり、あるノード（辺が $a$ 本接続されている）を選んで塗り替えると、赤と青を結ぶ辺の本数は $a - 2b$ 本変化する（増える）。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/196351/5b93fbee-e35a-8bf2-c4ba-436ddbf858c0.png)
今関心があるのは「赤と青を結ぶ辺の本数が奇数か偶数か」だけで、具体的な本数は必要ない。すると、不明な変数である $b$ がいくつかは無関係に、$a$ が奇数なら総本数の奇偶は反転し、$a$ が偶数なら変化しない。

ということでDと同様に動的プログラミングで、赤いノードの個数 $0 \leq x \leq k$ と赤と青を結ぶ辺の本数の奇偶 $0 \leq y \leq 1$ を追跡する…とやると、$N \leq 2 \times 10^5$ という規模から間に合わない。

ノードに繋がる辺の本数を調べることで、繋がる辺の本数が偶数であるノードの個数 $e$ と、奇数であるノードの個数 $o$（$e + o = N$）は数えられる。赤く塗るノードのうち、繋がる辺が奇数本であるノードを必ず偶数個（$i$ 個）選ぶことで総本数を偶数にできる。そのやり方は全部で
$$\sum\_{i=0}^k {}\_oC\_i \cdot {}\_eC\_{k-i} \ \  (i は i \leq o, k-i \leq e を満たす偶数のみ)$$
となる。
これは、${}\_nC\_k$を計算する関数を `c :: Int -> Int -> Int` とすると次のようにコードにできる。
```haskell
r $ sum [mul (c o i) (c e j) | i <- [0,2 .. k], let j = k - i, i <= o, j <= e]
```

大きな $n$ に対する ${}\_nC\_k$ の計算は結構大変で、上の式を計算するのはやはり無謀そうだが、[アルゴリズムロジック / 競プロでよく使う二項係数(nCk)を素数(p)で割った余りの計算と逆元のまとめ / §1.2](https://algo-logic.info/combination/#toc_id_1_2)「$n$ が巨大 ($k \leq 10^7, n \leq 10^9$) でPが素数のとき」に求め方が説明されている。

## 二項係数

以下、計算は全て$\bmod p$ で考える。（$N$ が大きいので `Data.Vector.Unboxed`, `Data.Vector.Unboxed.Mutable` を使用していく。）

$\displaystyle {}\_nC\_k = \frac{n !}{k ! \cdot (n-k) !}$ を求めたい。
$n$ の上限を定めておき、$1 \leq p \leq n$ に関して $p !$ と $(p !)^{-1}$ を計算してベクタに持っておく。それぞれ `fact, factinv :: UV.Vector Int` とすると `c` は上の式をそのまま書き下すことで作れる。

```haskell
c n k = (fact UV.! n) `mul` (factinv UV.! k) `mul` (factinv UV.! (n - k))
```

階乗 `fact` は、普通のリストの `scanl` で作っておいて `fromListN` でベクタ化すればよい。N付きの方が要素数が先に決まるため length しないぶん効率的な気がして選択。

```haskell
fact = UV.fromListN (succ n) $ scanl mul 1 [1 .. n]
```

これに逆数関数をmapすれば `factinv` は作れなくもないが、少し重い。
何やら魔法のようなやり方で、1からNまでの逆数列 `inv` を作ることができる方法を写経する。

```haskell
mkInvs :: Int -> UV.Vector Int
mkInvs n = UV.create (do
  v <- MUV.new (succ n)
  MUV.write v 0 1
  MUV.write v 1 1
  forM_ [2..n] (\a -> do
    let (q, r) = divMod p a
    rinv <- MUV.read v r
    let ainv = mul (negate q) rinv
    MUV.write v a ainv
    )
  return v
  )
```

これを前から掛けていけば `factinv` も作れる。

```haskell
factinv = UV.scanl1 mul $ mkInvs n
```

## 奇数辺ノードの数え上げ

$o$ を求めるところが残っている。`uvs` の中に各ノード番号が現れるたびに奇偶を反転し、最後に奇数となっているものを数えればよい。

```haskell
countOdds :: Int -> [(Int,Int)] -> ST s Int
countOdds n uvs = do
  vec <- MUV.replicate (succ n) False
  forM_ uvs (\(u,v) -> do
    MUV.modify vec not u
    MUV.modify vec not v
    )
--  MUV.foldl (\cnt b -> if b then succ cnt else cnt) 0 vec
  foldM (\cnt i -> do
    b <- MUV.read vec i
    return $ if b then succ cnt else cnt
    ) 0 [1 .. n]
```

Oops! AtCoderのVectorライブラリが微妙に古くて `foldl` がないため `foldM` に展開する必要があった。

## 結果

```haskell
import qualified Data.Vector.Unboxed.Mutable as MUV
import qualified Data.Vector.Unboxed as UV
import Control.Monad.ST

p = 998244353
r x = mod x p
mul x y = r (x * y)

abc262e :: Int -> Int -> Int -> [(Int,Int)] -> Int
abc262e n m k uvs = r $ sum
    [mul (c o i) (c e j) | i <- [0, 2 .. k], i <= o, let j = k - i, j <= e]
  where
    o = runST (countOdds n uvs)
    e = n - o
    c = mkComb n

countOdds :: Int -> [(Int,Int)] -> ST s Int
countOdds n uvs = do
  vec <- MUV.replicate (succ n) False
  forM_ uvs (\(u,v) -> MUV.modify vec not u >> MUV.modify vec not v)
  foldM (\cnt i -> do
    b <- MUV.read vec i
    return $ if b then succ cnt else cnt
    ) 0 [1 .. n]

mkComb :: Int -> (Int -> Int -> Int)
mkComb nmax = c
  where
    c n k = (fact UV.! n) `mul` (factinv UV.! k) `mul` (factinv UV.! (n - k))
    fact = UV.fromListN (succ nmax) $ scanl mul 1 [1 .. nmax]
    factinv = UV.scanl1 mul invs
    invs = UV.create (do
      v <- MUV.new (succ nmax)
      MUV.write v 0 1
      MUV.write v 1 1
      forM_ [2..nmax] (\a -> do
        let (q, r) = divMod p a
        rinv <- MUV.read v r
        let ainv = mul (negate q) rinv
        MUV.write v a ainv
        )
      return v
      )
```