PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: biggestotaku
Tester: pols_agyi_pols
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic programming, math
PROBLEM:
For a zero-indexed sequence S of length M, define f(S) = \prod_{i=0}^{M-1} (S_i + i)
Then define f(A, K) to be the sum of f(S) over all subsequences S of A that have length K.
You have an empty array B.
For i = 1, 2, 3, \ldots, N, do the following:
- Insert X_i at the start of B.
- Then, given K_i, compute \sum_{i=0}^{K_i} (K_i-i+1)^2 \cdot f(B, i)
EXPLANATION:
Recall that in the easy version of the problem, the parameter K (i.e. maximum allowed subsequence length) was fixed, which allowed us to maintain the overall answer by storing partial answers for subsequences.
In particular, we were not actually storing the answer for each possible subsequence length - we just combined everything since we only cared about the sum.
Here however, each query has a different K, so we can no longer combine everything.
It’s still possible to solve for each value of K separately, but that would lead to a complexity of \mathcal{O}(NK^2) which is too slow.
Instead, what we will do is explicitly maintain the f(B, i) values as we update B.
If we’re able to do this, then computing the answer for any query K is trivial in \mathcal{O}(K) time, which is fine complexity-wise.
For ease of notation, let us define f_m(i) to be the sum of products of all length-i subsequences, given that m elements have been inserted into B.
Suppose the m-th element inserted is X.
Then, to compute f_m(i),
- For subsequences that don’t include the newly inserted X, the sum of values is simply f_{m-1}(i).
- Subsequences that do include the new value can be obtained by taking a length (i-1) subsequence from the existing elements and inserting X at the beginning.
This however will come at the cost of incrementing the values of each term in the product by 1, so we need to be able to compute the new sum of products.
If you try working things out for a few small values of i, for example i = 1, 2, 3, you may notice that f_m(i) can be written as \sum_{j \le i} c_j f_{m-1}(j) for some appropriate coefficients c_j that depend only on m, i and j.
(This is not hard to do, just expand a couple of sums.)
This naturally leads to the guess that such a relation holds for every i, i.e. we can always find an appropriate set of coefficients dependent on m, i, j such that f_m(i) = \sum_{j \le i} c_j \cdot f_{m-1}(i).
If this were the case, and the coefficients can be computed quickly enough, then we’ll have something of a solution to the problem.
Indeed, it turns out that quite nicely, we have:
This can be proved using induction.
The above recurrence allows us to compute f_m(i) in \mathcal{O}(i) time if applied directly.
However, this can be optimized by noting that we can simply store cumulative sums of f_{m-1}(j) \cdot ((m-1-j)!) which allows us to compute all the values of f_m(i) for i \le K in \mathcal{O}(K) time.
Since K \le 1000 we can hence store the values of f_m(i) for all i \le 1000 this way, and for some query K simply output the required sum in \mathcal{O}(K).
TIME COMPLEXITY:
\mathcal{O}(N\cdot \max K) per testcase.
CODE:
Author's code (C++)
#include<bits/stdc++.h>
#define int long long
#define nl "\n"
using namespace std;
const int maxn = 1e5 + 10;
const int mod = 998244353;
int f[maxn], rf[maxn], iv[maxn];
void precalc() {
f[0] = f[1] = rf[0] = rf[1] = iv[1] = 1;
for (int i = 2; i < maxn; i++) {
f[i] = f[i - 1] * i % mod;
iv[i] = mod - mod / i * iv[mod % i] % mod;
rf[i] = rf[i - 1] * iv[i] % mod;
}
}
void solve(int cid) {
const int mod = 998244353;
int n;
cin >> n;
const int max_k = 1000;
vector<int> dp(max_k + 1);
dp[0] = 1;
for (int bef = 0; bef < n; bef++) {
int X;
cin >> X;
for (int i = 0, sm = 0; i <= max_k; i++) {
int p_val = dp[i];
dp[i] = (dp[i] + X * sm % mod * rf[bef - i + 1]) % mod;
if (bef < i) break;
sm = (sm + p_val * f[bef - i]) % mod;
}
int K;
cin >> K;
int res = 0;
for (int i = 0; i <= K; i++) {
res = (res + (K + 1 - i) * (K + 1 - i) * dp[i]) % mod;
}
cout << res << nl;
}
}
signed main() {
ios_base::sync_with_stdio(0);
cin.tie(0);cout.tie(0);
precalc();
int t = 1;
cin >> t;
int cid = 0;
while(t--){
solve(++cid);
}
}
}