k乗和と exp

一般に，k乗和を見たら \(e^{ax}\) の総積を考えると上手くいくことがあります．
これは以下の性質が成り立つからです．

• \(e^{ax} e^{bx} = e^{(a+b)x}\)（性質1）
• \(k![x^k]e^{ax} = a^k\)（性質2）

ただし，多項式 \(f(x)\) の k 次の係数を \([x^k]f(x)\) で表しています．

解法

今回の問題は

\[\sum_{l,r}(\sum_{i\in[l,r]} A_i)^k\]

ですが，これは難しいので代わりにk乗和を総積に置き換えた以下のような問題を考えてみます．

\[\sum_{l,r}\prod_{i\in[l,r]} f_i\]

これは各 \(r\) について \(\sum_{l}\prod_{i\in[l,r]} f_i\) を足し合わせればよく，以下のようなコードで線形時間で解くことが出来ます．

def solve(f: list[int]) -> int:
    result = 0
    sum = 0
    for x in f:
        sum += 1
        sum *= x
        result += sum
    return result

ここで \(f_i \coloneqq e^{A_ix}\) を代入してみると（性質1）より

\[ \sum_{l,r}\prod_{i\in[l,r]} f_i = \sum_{l,r}\prod_{i\in[l,r]} e^{A_ix} = \sum_{l,r}e^{\sum_{i\in[l,r]} A_ix} \]

が成り立ちます．
したがって（性質2）より \(k![x^k]\sum_{l,r}\prod_{i\in[l,r]} f_i\) が今回の問題の答えです．

ここで FPS は全て \(k\) 次の項まで持てば充分であるため，足し算は \(\Theta(k)\)，掛け算は \(\Theta(k^2)\) 時間で行えることに留意すると，全体で \(\Theta(nk^2)\) 時間程度でこの問題を解くことが出来ます．

実装例

#include <atcoder/all>
#include <bits/stdc++.h>

using mint = atcoder::modint998244353;
using poly = std::vector<mint>;

int k;
poly poly_exp(int a) {
    poly res(k + 1, 1);
    for (int i = 1; i <= k; i++)
        res[i] = res[i - 1] * mint::raw(a) / mint::raw(i);
    return res;
}
void multiply(poly &f, const poly &g) {
    f = atcoder::convolution(f, g);
    f.resize(k + 1);
}

int main() {
    int n;
    std::cin >> n >> k;

    mint ans = 0;
    poly sum(k + 1, 0);

    while (n--) {
        int a;
        std::cin >> a;
        sum[0]++;
        multiply(sum, poly_exp(a));
        ans += sum[k];
    }

    for (int i = 1; i <= k; i++)
        ans *= i;

    std::cout << ans.val() << std::endl;
}

似たような話 https://noshi91.hatenablog.com/entry/2022/12/31/015729