EDPC M - CandiesをFPSで

Last updated Apr 13, 2022

$$ \begin{align} & \prod_{i = 1}^{N} (1 + x + \cdots + x^{a_i}) \\\ & = \prod_{i = 1}^{N} \frac{1 - x^{a_i + 1}}{1 - x} \\\ & = (1 - x)^{-N} \prod_{i = 1}^{N} (1 - x^{a_i + 1}) \end{align} $$

$\prod_{i = 1}^{N} (1 - x^{a_i + 1})$はdpで計算して$O(N K)$、$(1 - x)^{-N}$はN回累積和の操作に対応するので$O(N K)$です。

提出コード

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
i64 N, K;
cin >> N >> K;
vector<i64> A(N);
rep(i,0,N) {
  cin >> A[i];
  A[i]++;
}
vector<fp> dp(K + 1);
dp[0] = 1;
rep(i,0,N) {
  for(i64 j = K + 1; j --> A[i];) {
    dp[j] -= dp[j - A[i]];
  }
}
rep(i,0,N) {
  rep(j,1,K + 1) {
    dp[j] += dp[j - 1];
  }
}
cout << dp[K] << endl;

負の二項定理を使えば、最後のパートを$(N + K)$で計算できそう。

提出コード

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
i64 N, K;
cin >> N >> K;
vector<i64> A(N);
rep(i,0,N) {
  cin >> A[i];
  A[i]++;
}
vector<fp> dp(K + 1);
dp[0] = 1;
rep(i,0,N) {
  for(i64 j = K + 1; j --> A[i];) {
    dp[j] -= dp[j - A[i]];
  }
}
fact.build(N + K);
fp ans;
rep(i,0,K + 1) {
  ans += dp[i] * fact.binom(K - i + N - 1, N - 1);
}
cout << ans << endl;