F - Range Power Sum Editorial by en_translator
Rephrase the problem statement as follows.
- There are \(N\) boxes arranged from left to right. Box \(i\) contains \(A_i\) balls.
- How many ways are there to perform the following procedure?
- Place two indistinguishable partitions between boxes in the sequence.
- For each \(i=1,2,\dots,K\), choose one ball in a box between the two partitions, and put a sticker labeled \(i\). A single ball may have multiple stickers.
Let us inspect the ball from the left, and determine the position of partitions while putting stickers. We will perform the following DP.
- \(\mathrm{dp}[i][j][k]\): the number of ways to inspect the first \(i\) boxes, while inserting \(j\ (0\leq j\leq 2)\) partitions and putting \(k\ (0\leq k\leq K)\) stickers
The transitions are as follows.
- \(\mathrm{dp}[i][j+1][k]\) += \(\mathrm{dp}[i][j][k]\)
- \(\mathrm{dp}[i+1][j][k]\) += \(\mathrm{dp}[i][j][k]\)
- \(\mathrm{dp}[i+1][1][k+p]\) += \(\mathrm{dp}[i][1][k] \times \binom{K-k}{p} \times A_i^p\ \ \ (1 \leq p\leq K-k)\)
The first transition corresponds to placing a partition, and the second to advancing to the next box without putting a sticker. The third corresponds to putting a sticker on the ball in the box currently inspecting, where
- the variable \(c\) corresponds to the number of stickers to put,
- the coefficient \(\binom{K-k}{p}\) to the number of ways to choose \(p\) stickers out of \(K-k\) stickers not attached yet, and
- the coefficient \(A_i^p\) to the number of ways to choose balls to put sticker on for each of the \(p\) chosen sticker.
Note that the third transition is only limited to \(j=1\) because one can put a sticker on a ball only when it is in a box between the partitions.
The tie complexity is \(O(NK^2)\).
Sample code (C++):
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using mint = atcoder::modint998244353;
int main() {
int n, k;
cin >> n >> k;
vector<int> a(n);
for (int &i: a) cin >> i;
vector binom(k + 1, vector<mint>(k + 1));
for (int i = 0; i <= k; i++) {
binom[i][0] = 1;
for (int j = 1; j <= i; j++) {
binom[i][j] = binom[i - 1][j] + binom[i - 1][j - 1];
}
}
vector dp(n + 1, vector<mint>(k + 1));
mint ans = 0;
for (int i = 0; i < n; i++) {
vector<mint> pw(k + 1); // pw[j] = (a[i])^j
pw[0] = 1;
for (int j = 0; j < k; j++) {
pw[j + 1] = pw[j] * a[i];
}
dp[i][0] += 1;
for (int j = 0; j <= k; j++) {
for (int l = 0; j + l <= k; l++) {
dp[i + 1][j + l] += dp[i][j] * binom[k - j][l] * pw[l];
}
}
ans += dp[i + 1][k];
}
cout << ans.val() << endl;
}
Note that in the sample code above, the second key \(j\) of the DP is eliminated some how. Consider how it is eliminated. (Hint: line 28 and 34)
posted:
last update: