AGC 025 B - RGB Coloring

• 問題

問題概要

ブロックに0,A,B,A+Bのいずれかの得点をつけることができる
合計Kの得点をつける塗り方をmod 998244353で求めよ.
1 <= N,A,B <= 310⁵, 0 <= K <= 1810¹⁰

解法

単純にDPを立てると,状態がO(NK)とかになってしまうので間に合わない.
ここで問題特有の3つ目の得点がA+Bであると言うことに注目してみよう.
A+Bをつけると言うのを,AとBのどちらの得点もつけると考えることができる.
すなわちAとBの得点を独立につけることができると言うことである.

以上のことから,aA+bB=K(1<=a,b<=N) を満たすaを全探索し,
nCa * nCb を足して行くことで O(N)で解ける.

メモ

f(N,K) = f(N-1,K) + f(N-1,K-A) + f(N-1,K-B) + f(N-1,K-A-B) で答えを求められることがわかってそれをなんとか高速化しようとしたら終わってしまった.
問題特有の要素に注目するの大事.

コード