ABC262 A～E をHaskellで

A - World Cup

問題 ABC262A

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

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

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

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

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

結果

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

B - Triangle (Easier)

問題 ABC262B

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

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

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

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

これが動くようにするには、ノード番号をキーとして、接続されている（より大きい）ノード番号の一覧を取得できるようにする必要がある。毎回 uvs をスキャンして

[v | (u,v) <- uvs, a == u] -- aと結ばれている頂点

などとするのは計算機に酷なので、そのような配列をData.Array.accumArray で先に構築しておくとよい。

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

結果

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

C - Min Max Pair

問題 ABC262C

シグネチャを決める。

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$ となる。

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 とリストのままするのは遅いので、配列に移しておくべきである。

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

結果

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]

D - I Hate Non-integer Number

問題 ABC262D

シグネチャを決める。

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

モジュロ演算の足し算を作っておく。

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 を用いる。

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]$ が求める値である。

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)$ で、上述の手順を走らせるよりも速く結果を求められるが、そこを凝る必要もないだろう。

結果

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

シグネチャを決める。

abc262e :: Int -> Int -> Int   -- N, M, K
        -> [ (Int,Int) ]       -- Ui, Vi のベアのリスト uvs
        -> Int                 -- 答え

モジュラな演算なども定義しておく。

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

ノードは最初全て青く塗られているとし、そのうちいくつか選んで赤く塗り替えていくことを考える。下の図で示すとおり、あるノード（辺が $a$ 本接続されている）を選んで塗り替えると、赤と青を結ぶ辺の本数は $a - 2b$ 本変化する（増える）。

今関心があるのは「赤と青を結ぶ辺の本数が奇数か偶数か」だけで、具体的な本数は必要ない。すると、不明な変数である $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 とすると次のようにコードにできる。

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「$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 は上の式をそのまま書き下すことで作れる。

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

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

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

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

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 も作れる。

factinv = UV.scanl1 mul $ mkInvs n

奇数辺ノードの数え上げ

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

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 に展開する必要があった。

結果

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
      )