Official

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?
    1. Place two indistinguishable partitions between boxes in the sequence.
    2. 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: