AtCoderRegularContest134 C問題500点「The Majority」
https://gyazo.com/83df7a2ec990e4b4f64f261c77fc730d
問題概要
制約
$ N \leq 10^5
$ K\leq200
解法・お気持ち
基本的には公式解説の通りです。
ただし、解説が非常にシンプルで読み解くのが難しかったため、どのような気持ちでその計算式になっているかまとめます。
それでは、公式解説の計算式についてまとめます。
$ \binom{a_1 - 1 - \sum_{i=2}^{N}a_i}{K-1} \prod_{i=2}^{N}\binom{a_i + K-1}{K-1}
上記の計算式についてです。
まず、右側の$ \prod_{i=2}^{N}\binom{a_i + K-1}{K-1}についてです。
この項はそもそも i = 2 ~ N のボールをそれぞれ K 個の箱に入れる組み合わせを表しています。
そして、$ \binom{a_i + K - 1}{K-1}は、ボール$ iを$ K個の箱に入れるときの組み合わせである重複組み合わせを表しています。 たとえば、$ i = 2のとき、以下のように $ K -1本の仕切りで２のボールを分ける、というように考えられます。
https://gyazo.com/ff96b54cf96d49614ff7a8ef5cdf2ab7
$ i = 2, i=3, i = \dotsなどはそれぞれ独立して考えられるため、積の法則で$ \prod が使われています。
次に、$ \binom{a_1 - 1 - \sum_{i=2}^{N}a_i}{K-1}についてです。
この項は 自由に配置できる D 個の 1 のボールを K 個の箱に入れる組み合わせを表しています。
ここで指している D とは、$ K個のハコに1 個ずつ入れる用のK個、そして $ \sum_{i=2}^N a_i個のボールのための相殺分を除いたものです。（つまり、$ D = a_1 - K - \sum_{i=2}^N a_i）
これによって、 D 個のボールは自由に $ K個のハコに入れられるようになります。
ただし、それだと以上の式になりません。これは式変形された状態だからです。
箱に1個ずつ入れておく&i=2~Nのボールの相殺を除いた、自由に配置できるボールは $ a_1 - K - \sum_{i=2}^N a_i でした。
これら自由に配置できるボールを同様に重複組み合わせで$ K個の箱に入れようとすると
$ \binom{a_1 - K - \sum_{i=2}^N a_i +(K - 1)}{K-1} = \binom{a_1 - 1 - \sum_{i=2}^{N}a_i}{K-1}
となり、公式解説通りの結果になりました。
よって、あとはこの公式解説通りの計算式を実装すればよいです。
追加の留意点として、組合せの計算方法に注意が必要です。 分子に来る値が$ a_1, \sum_iなど非常に大きな値になります。
よって、あらかじめ mod を計算しておくことはできません。（よくある二項計算の前計算を指しています。）
しかし、分子は$ K-1でありたかだか 200 しかありません。
そこで、実際に紙で行うような組合せの計算を行えばいいです。 （ソースコードを参考にしてみてください。）
計算量
$ O(NK)
新たな学び
過半数（他＋１）では、先に1個ずつおいておくなどができる
計算しやすい重複組み合わせのような状態に持っていく事が大事
反省点
inv の計算はあらかじめしておくと良い
mod の割り算は重いので（繰り返し自乗法などでさらに logN かかる）
コード
code: c++
constexpr int mod = 998244353;
constexpr int MX = 2000200;
using MD = ModInt<mod>;
int main() {
int N, K;
cin >> N >> K;
vector<ll> A(N);
cin >> A;
vector<MD> inv(K + 1);
for (int i = 0; i <= K; i++) invi = MD(1) / MD(i); ll B = 0;
for (int i = 1; i < N; i++) B += Ai; auto C = &(long long x, int y) -> MD { if (x < 0) return MD(0);
if (y < 0 || y > x) return MD(0);
auto ret = MD(1);
for (int i = 0; i < y; i++) {
ret = ret * (x - i) * invi + 1; }
return ret;
};
MD ans = 1;
// 1 の置き方
ans *= C(A0 - B - 1, K - 1); // 2 ~ N の置き方
for (int i = 1; i < N; i++) {
ans *= C(Ai + K - 1, K - 1); }
cout << ans << endl;
}