A - World Cup
まず、関数の型を次のように決める。
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 :: 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 :: 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 :: 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 :: 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
)