题目描述

给定正整数 N, K 和一个长度为 N 的整数序列 A = (A₁, A₂, …, A<0xE2><0x82><0x99>).

计算以下表达式的值，结果对 998244353 取模：
$$\sum_{1 \le l \le r \le N} \left( \sum_{l \le i \le r} A_i \right)^K$$
即，对所有可能的连续子数组 A[l...r]，计算其元素和的 K 次方，并将这些 K 次方的值全部加起来。

样例

输入:
3 2
3 1 2

输出:
75

说明：
子数组和及其平方为：
[3]: 3² = 9
[1]: 1² = 1
[2]: 2² = 4
[3, 1]: (3+1)² = 4² = 16
[1, 2]: (1+2)² = 3² = 9
[3, 1, 2]: (3+1+2)² = 6² = 36
总和 = 9 + 1 + 4 + 16 + 9 + 36 = 75.

输入:
10 5
91 59 85 60 57 72 12 3 27 16

输出:
428633385

注意结果需要取模 998244353。

算法 (动态规划 / 贡献法)

$O(N K^2)$ 或 $O(N K \log K)$

思路分析

1. 问题转化与前缀和:
   直接枚举所有 l 和 r (共 O(N²) 对) 并计算子数组和 (O(N) 或 O(1) 如果用前缀和) 再求 K 次方 (O(log K))，总复杂度为 O(N³ log K) 或 O(N² log K)，对于 N=2e5 来说太慢了。
   令 P[i] 表示序列 A 的前缀和，即 P[i] = A₁ + A₂ + ... + Aᵢ，并定义 P[0] = 0。
   那么子数组 A[l...r] 的和为 P[r] - P[l-1]。
   我们需要计算的式子变为：
   $$\sum_{r=1}^{N} \sum_{l=1}^{r} (P[r] - P[l-1])^K \pmod{998244353}$$

2. 二项式展开:
   由于 K 很小 (K ≤ 10)，我们可以利用二项式定理展开 (P[r] - P[l-1])^K:
   $$(P[r] - P[l-1])^K = \sum_{k=0}^{K} \binom{K}{k} (P[r])^{K-k} (-P[l-1])^k$$
   $$= \sum_{k=0}^{K} \binom{K}{k} (-1)^k (P[r])^{K-k} (P[l-1])^k$$
   其中 binom(K, k) 是组合数 C(K, k)。

3. 交换求和顺序:
   将展开式代入原求和式：
   $$\text{Ans} = \sum_{r=1}^{N} \sum_{l=1}^{r} \sum_{k=0}^{K} \binom{K}{k} (-1)^k (P[r])^{K-k} (P[l-1])^k$$
   交换求和顺序，先把关于 k 的求和提到外面：
   $$\text{Ans} = \sum_{k=0}^{K} \binom{K}{k} (-1)^k \sum_{r=1}^{N} \sum_{l=1}^{r} (P[r])^{K-k} (P[l-1])^k$$
   这个形式仍然难以直接计算内部的双重求和。

4. 考虑贡献 / 动态规划:
   我们尝试按 r 从 1 到 N 的顺序进行计算。当我们处理到 r 时，我们希望能够快速计算出所有以 r 结尾的区间的贡献，即 ∑_{l=1}^{r} (P[r] - P[l-1])^K。
   令 Term(r) = ∑_{l=1}^{r} (P[r] - P[l-1])^K。那么 Ans = ∑_{r=1}^{N} Term(r)。
   将二项式展开代入 Term(r)：
   $$Term(r) = \sum_{l=1}^{r} \sum_{k=0}^{K} \binom{K}{k} (-1)^k (P[r])^{K-k} (P[l-1])^k$$
   再次交换求和顺序：
   $$Term(r) = \sum_{k=0}^{K} \binom{K}{k} (-1)^k (P[r])^{K-k} \left( \sum_{l=1}^{r} (P[l-1])^k \right)$$
   令 S[k][i] = ∑_{j=0}^{i} (P[j])^k，即前缀和的 k 次方的和。
   那么 ∑_{l=1}^{r} (P[l-1])^k = ∑_{j=0}^{r-1} (P[j])^k = S[k][r-1]。
   所以：
   $$Term(r) = \sum_{k=0}^{K} \binom{K}{k} (-1)^k (P[r])^{K-k} S[k][r-1]$$
   这个 Term(r) 的计算看起来是可行的。我们需要在计算 Term(r) 时，能够快速得到 S[k][r-1] 的值。

5. 维护 S[k]:
   注意到 S[k][r] = S[k][r-1] + (P[r])^k。
   这启发我们可以用一个大小为 K+1 的数组（我们称之为 SumPow，对应代码中的 S 数组）来维护当前的 S[k][r-1] 的值。
   当我们从 r-1 移动到 r 时：

   • 首先，利用当前的 SumPow 数组（它存储的是 S[k][r-1] for k=0..K）和当前的 P[r] 来计算 Term(r)，并累加到最终答案 ans 中。
   • 然后，更新 SumPow 数组，使其存储 S[k][r] 的值，以备计算 Term(r+1) 时使用。更新方式为：SumPow[k] = SumPow[k] + (P[r])^k for k=0..K。

6. 算法流程:

   • 预计算组合数 binom(K, k) 模 998244353。
   • 初始化前缀和 P = 0。
   • 初始化 SumPow 数组（大小为 K+1），SumPow[0] = 1，SumPow[k] = 0 for k > 0。（这对应于 S[k][-1]，其中 P[-1]=0，所以 P[-1]^0=1, P[-1]^k=0 for k>0）。
   • 初始化最终答案 ans = 0。
   • 遍历 r 从 0 到 N-1 (使用 0-based 索引处理数组 A):
     1. 更新当前前缀和 P = (P + A[r]) % mod。
     2. 计算 Term(r) 的贡献：
        current_term_sum = 0
for k = 0 to K:
sign = (k % 2 == 0) ? 1 : -1
term = binom(K, k) * sign * power(P, K - k) * SumPow[k]
current_term_sum = (current_term_sum + term) % mod
(需要处理好负数取模)
     3. 将 current_term_sum 加入 ans: ans = (ans + current_term_sum) % mod。
     4. 更新 SumPow 数组：
        for m = 0 to K:
SumPow[m] = (SumPow[m] + power(P, m)) % mod
   • 输出 ans。

7. 复杂度分析:

   • 预计算组合数：O(K)。
   • 主循环 N 次。
   • 在循环内部：
     • 计算 Term(r) 需要一个循环 K+1 次。内部涉及组合数查找 (O(1))、模幂 power(P, K-k) (O(log K) 或 O(K) 如果预计算 P 的幂)、乘法和加法。总共 O(K log K) 或 O(K^2)。
     • 更新 SumPow 需要一个循环 K+1 次。内部涉及模幂 power(P, m) (O(log K) 或 O(K)) 和加法。总共 O(K log K) 或 O(K^2)。
   • 总时间复杂度为 O(N * (K log K)) 或 O(N * K^2)。考虑到 K=10，这两种复杂度都足够快。O(N * K^2) 约为 2e5 * 100 = 2e7，在 2 秒内是可行的。

8. 代码实现细节:

   • 代码使用了 ModInt 模板类来处理模运算，确保所有中间结果都在模 M 下。
   • Comb 结构体用于预计算和查询组合数。
   • power 函数用于计算模幂。
   • P 变量对应思路中的 P[r]。
   • S vector 对应思路中的 SumPow 数组。
   • 代码中的 r 循环从 0 到 N-1，P 对应 P[r] (使用 0-based A 计算的前缀和)，S[k] 对应 S[k][r-1]（即 ∑_{j=0}^{r-1} (P[j])^k）。
   • 内层第一个 k 循环计算的就是 Term(r)。
   • 内层第二个 m 循环更新 S 数组（SumPow）为 S[m][r]。
   • 逻辑与推导一致。

时间复杂度

O(N * K^2) 或 O(N * K log K)，取决于模幂的计算方式。

C++ 代码

#include <bits/stdc++.h> // 引入所有标准库

using namespace std;
using i64 = long long; // 定义 i64 为 long long 的别名

// ... [省略了 ModInt, Barrett, Comb 等模板代码] ...
// 这些模板提供了模运算、组合数计算等基础功能
// Z 是 ModInt<998244353> 的别名

// Comb 结构体用于计算组合数
struct Comb {
    int n;
    std::vector<Z> _fac; // 阶乘
    std::vector<Z> _invfac; // 阶乘的逆元
    std::vector<Z> _inv; // 逆元

    Comb() : n{0}, _fac{1}, _invfac{1}, _inv{0} {}
    Comb(int n) : Comb() {
        init(n); // 初始化到 n
    }

    // 初始化组合数计算所需的表，直到 m
    void init(int m) {
        if (m <= n) return;
        _fac.resize(m + 1);
        _invfac.resize(m + 1);
        _inv.resize(m + 1);

        for (int i = n + 1; i <= m; i++) {
            _fac[i] = _fac[i - 1] * i;
        }
        _invfac[m] = _fac[m].inv(); // 计算最大阶乘的逆元
        for (int i = m; i > n; i--) { // 反向计算逆元和阶乘逆元
            _invfac[i - 1] = _invfac[i] * i;
            _inv[i] = _invfac[i] * _fac[i - 1];
        }
        n = m;
    }

    // 返回 m!
    Z fac(int m) {
        if (m > n) init(2 * m); // 动态扩展
        return _fac[m];
    }
    // 返回 1/m!
    Z invfac(int m) {
        if (m > n) init(2 * m);
        return _invfac[m];
    }
    // 返回 1/m
    Z inv(int m) {
        if (m > n) init(2 * m);
        return _inv[m];
    }
    // 计算 C(n, m) = n! / (m! * (n-m)!)
    Z binom(int n, int m) {
        if (n < m || m < 0) return 0; // 处理边界情况
        return fac(n) * invfac(m) * invfac(n - m);
    }
} comb; // 全局组合数对象

// 主函数
int main() {
    // 加速 IO
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, K; // 输入 N 和 K
    cin >> N >> K;
    vector<Z> A(N); // 存储序列 A (使用模整数类型 Z)
    for (int i = 0; i < N; i ++ ) {
        cin >> A[i]; // 读入序列 A
    }

    Z P = 0;  // 当前前缀和 P[r]
    // S[k] 存储 sum_{j=0}^{r-1} (P[j])^k
    vector<Z> S(K + 1, 0);  
    S[0] = 1; // P[-1]^0 = 0^0 = 1 (边界情况)
    Z ans = 0; // 最终答案

    // 遍历区间的右端点 r (从 0 到 N-1)
    for (int r = 0; r < N; r ++ ) {
        // 更新当前前缀和 P = P[r]
        P += A[r];  

        // 计算以 r 结尾的所有区间的贡献: sum_{l=0}^{r-1} (P[r] - P[l])^K
        // 注意这里的 l 对应 P[l-1] in 1-based index， P[l] in 0-based index
        // for (int l_minus_1 = -1; l_minus_1 < r; ++l_minus_1)
        // (P[r] - P[l_minus_1])^K
        for (int k = 0; k <= K; k ++ ) {
            // 应用展开式: C(K, k) * (-1)^k * P[r]^(K-k) * S[k]
            // S[k] 存储的是 sum_{j=0}^{r-1} P[j]^k
            Z sign = (k % 2 == 0) ? Z(1) : Z(-1); // 计算 (-1)^k
            // 计算当前项贡献
            Z term = comb.binom(K, k) * sign * power(P, K - k) * S[k];
            // 累加到最终答案
            ans += term;
        }

        // 更新 S 数组，为下一次迭代 (r+1) 做准备
        // S[m] 需要从 sum_{j=0}^{r-1} P[j]^m 更新为 sum_{j=0}^{r} P[j]^m
        for (int m = 0; m <= K; m ++ ) {
            // S[m] = S[m] + P[r]^m
            S[m] += power(P, m);
        }
    }

    // 输出最终答案
    cout << ans << "\n";

    return 0;
}