[科技] 求数列的前k次方和
[科技] 求数列的前k次方和
到现在才会的一个科技,写一篇博客来记录一下。
简单来说,就是对于\(0 \leq t \leq k\)求\(\sum_{i = 1} ^ n a_i^t\),\(n, k \leq 10^5\)。
我们考虑答案序列的生成函数:
\[F(x) = \sum_{t = 0} ^ {\infty} x^t \sum_{i = 1} ^ n a_i ^ t = \sum_{i = 0} ^ {n}\sum_{t = 0} ^ {\infty} (a_ix) ^ t = \sum_{i = 1} ^ n \frac{1}{1 - a_ix}
\]
发现这样还是无法下手,但是考虑到$$(\ln(1 - a_ix))' = \ln'(1 - a_ix) \times a_i = \frac{-a_i}{1 - a_ix} = -\sum_{t = 0} ^ {\infty}(a_ix) ^ t a_i$$
我们记$$G(x) = \sum_{i = 1} ^ n (\ln(1 - a_ix))' = \sum_{i = 1} ^ n \frac{-a_i}{1 - a_ix} = -\sum_{i = 1} ^ {n} \sum_{t = 0} ^ {\infty}(a_ix) ^ t a_i = -\sum_{t = 0} ^ {\infty}x ^ t\sum_{i = 1} ^ n a_i ^ {t + 1}$$
那么\(F(x) = -xG(x) + n\),于是我们求出了\(G(x)\),就可以快速求出\(F(x)\)了。
考虑继续化简\(G(x)\):
\[G(x) = (\sum_{i = 1} ^ n \ln(1 - a_ix))' = \ln'(\prod_{i = 1} ^ n (1 - a_ix))'
\]
里面的那个分治\(FFT\)一下就行了,然后求\(Ln\),求导就可以得到\(G\)了。
不过不知道是不是自己的写法有问题,发现这样只能求到前\(n\)项的幂和,必须在\(a\)后面补零才行……
#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 50;
const int Md = 998244353;
typedef long long ll;
typedef vector<int> Vec;
inline int Add(const int &x, const int &y) { return (x + y >= Md) ? (x + y - Md) : (x + y); }
inline int Sub(const int &x, const int &y) { return (x - y < 0) ? (x - y + Md) : (x - y); }
inline int Mul(const int &x, const int &y) { return (ll)x * y % Md; }
int Powe(int x, int y) {
int ans = 1;
while(y) {
if(y & 1) ans = Mul(ans, x);
x = Mul(x, x);
y >>= 1;
}
return ans;
}
int n, k;
int a[N];
namespace Poly {
int rev[N << 2 | 1], inv[N << 2 | 1];
void Init() {
inv[0] = inv[1] = 1;
for(int i = 2; i < N; i++) {
inv[i] = Mul(Md - Md / i, inv[Md % i]);
}
}
void DFT(Vec &A, int len) {
for(int i = 0; i < len; i++) if(i < rev[i]) swap(A[i], A[rev[i]]);
for(int i = 1; i < len; i <<= 1) {
int wn = Powe(3, (Md - 1) / (i << 1));
for(int j = 0; j < len; j += i << 1) {
int nw = 1, x, y;
for(int k = 0; k < i; k++, nw = Mul(nw, wn)) {
x = A[j + k], y = Mul(nw, A[i + j + k]);
A[j + k] = Add(x, y); A[i + j + k] = Sub(x, y);
}
}
}
}
void IDFT(Vec &A, int len) {
reverse(A.begin() + 1, A.end());
int IV = Powe(len, Md - 2);
DFT(A, len);
for(int i = 0; i < len; i++) A[i] = Mul(A[i], IV);
}
Vec MUL(Vec A, Vec B) {
int n = A.size(), m = B.size(), len;
for(len = 1; len < n + m - 1; len <<= 1);
for(int i = 0; i < len; i++) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) ? len >> 1 : 0);
A.resize(len); B.resize(len);
DFT(A, len); DFT(B, len);
for(int i = 0; i < len; i++) A[i] = Mul(A[i], B[i]);
IDFT(A, len);
A.resize(n + m - 1);
return A;
}
Vec GetInv(Vec A, int len) {
Vec B(1, Powe(A[0], Md - 2)), C;
for(int i = 2; (i >> 1) < len; i <<= 1) {
for(int j = 0; j < (i << 1); j++) rev[j] = (rev[j >> 1] >> 1) | ((j & 1) ? i : 0);
C = A; C.resize(i);
C.resize(i << 1); DFT(C, i << 1);
B.resize(i << 1); DFT(B, i << 1);
for(int j = 0; j < (i << 1); j++) B[j] = Mul(B[j], Sub(2, Mul(B[j], C[j])));
IDFT(B, i << 1);
B.resize(i);
}
B.resize(len);
return B;
}
Vec Dir(Vec A) {
Vec B; int len = A.size(); B.resize(len);
for(int i = 1; i < len; i++) B[i - 1] = Mul(i, A[i]);
B[len - 1] = 0;
return B;
}
Vec Inter(Vec A) {
Vec B; int len = A.size(); B.resize(len);
for(int i = 1; i < len; i++) B[i] = Mul(A[i - 1], inv[i]);
B[0] = 0;
return B;
}
Vec Ln(Vec A, int len) {
A = Inter(MUL(Dir(A), GetInv(A, len)));
A.resize(len);
return A;
}
}
Vec Solve(int l, int r) {
if(l == r) {
if(l <= n) {
Vec A(2); A[0] = 1; A[1] = Md - a[l];
return A;
}
else {
Vec A(2); A[0] = 1; A[1] = 0;
return A;
}
}
int mid = (l + r) >> 1;
Vec tmp1 = Solve(l, mid), tmp2 = Solve(mid + 1, r);
tmp1 = Poly::MUL(tmp1, tmp2);
return tmp1;
}
int main() {
Poly::Init();
scanf("%d%d", &n, &k);
for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
Vec A = Solve(1, k);
A = Poly::Ln(A, A.size());
A = Poly::Dir(A);
for(int i = 0; i < A.size(); i++) A[i] = (Md - A[i]) % Md;
A.resize(A.size() + 1);
for(int i = A.size(); i; i--) A[i] = A[i - 1]; A[0] = 0;
A[0] = Add(A[0], n);
return 0;
}
「我不敢下苦功琢磨自己,怕终于知道自己并非珠玉;然而心中既存着一丝希冀,便又不肯甘心与瓦砾为伍。」