首页 > 基础资料 博客日记

线段树维护区间 k 次方和

2026-06-28 22:30:02基础资料围观18次

文章线段树维护区间 k 次方和分享给大家，欢迎收藏极客资料网，专注分享技术知识

维护一个长度为 \(n\) 的数组，支持以下操作：

  1. 区间赋值：将 \([l,r]\) 全部赋值为 \(x\)。
  2. 区间乘法：将 \([l,r]\) 全部乘以 \(c\)。
  3. 区间加法：将 \([l,r]\) 全部加上 \(d\)。
  4. 区间查询：查询 \(\sum_{i=l}^{r} a_i^k\)（\(k\) 次方和）。

记 \(S_t = \sum a_i^t\) 表示区间内所有数的 \(t\) 次方和，我们需要维护 \(t = 0,1,\cdots,K\) 的所有 \(S_t\)。

区间加

对于每个数 \(a_i\)，加上 \(d\) 后：

\[(a_i + d)^t = \sum_{j=0}^{t} \binom{t}{j} a_i^j d^{t-j} \]

对区间求和，记 \(S_t = \sum a_i^t\)：

\[S_t' = \sum_{j=0}^{t} \binom{t}{j} d^{t-j} S_j \]

区间乘

对于每个数 \(a_i\)，乘以 \(c\) 后：

\[(c \cdot a_i)^t = c^t \cdot a_i^t \]

对区间求和：

\[S_t' = c^t \cdot S_t \]

区间赋值

区间全部变为 \(x\)，长度为 \(len\)：

\[S_t = len \cdot x^t \]

朴素实现

时间复杂度为 \(O(n\log n+qk^2\log n)\)。

代码

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N=1e5+5,MAXK=100,P=998244353;
int n,m,a[N],C[MAXK+1][MAXK+1];
struct Node
{
    int lz1,lz2,lz3,len,s[MAXK+1];
    // lz1:赋值 -1表示无, lz2:乘法 1表示无, lz3:加法 0表示无
}tr[N<<2];

void init()
{
    for(int i=0;i<=MAXK;i++)
    {
        C[i][0]=C[i][i]=1;
        for(int j=1;j<i;j++)C[i][j]=(C[i-1][j-1]+C[i-1][j])%P;
    }
}

void push_assign(int u,int x)
{
    tr[u].lz1=x;
    tr[u].lz2=1;
    tr[u].lz3=0;
    int pw=1;
    for(int t=0;t<=MAXK;t++)
    {
        tr[u].s[t]=1LL*tr[u].len*pw%P;
        pw=1LL*pw*x%P;
    }
}

void push_mul(int u,int c)
{
    tr[u].lz2=1LL*tr[u].lz2*c%P;
    tr[u].lz3=1LL*tr[u].lz3*c%P;
    int pw=1;
    for(int t=0;t<=MAXK;t++)
    {
        tr[u].s[t]=1LL*tr[u].s[t]*pw%P;
        pw=1LL*pw*c%P;
    }
}

void push_add(int u,int d)
{
    int pd[MAXK+1]={1};
    for(int i=1;i<=MAXK;i++)pd[i]=1LL*pd[i-1]*d%P;
    int ns[MAXK+1]={0};
    for(int t=0;t<=MAXK;t++)
        for(int j=0;j<=t;j++)
            ns[t]=(ns[t]+1LL*C[t][j]*pd[t-j]%P*tr[u].s[j])%P;
    for(int t=0;t<=MAXK;t++)tr[u].s[t]=ns[t];
    tr[u].lz3=(tr[u].lz3+d)%P;
}

void pd(int u)
{
    if(tr[u].lz1!=-1)
    {
        push_assign(u<<1,tr[u].lz1);
        push_assign(u<<1|1,tr[u].lz1);
        tr[u].lz1=-1;
    }
    if(tr[u].lz2!=1)
    {
        push_mul(u<<1,tr[u].lz2);
        push_mul(u<<1|1,tr[u].lz2);
        tr[u].lz2=1;
    }
    if(tr[u].lz3)
    {
        push_add(u<<1,tr[u].lz3);
        push_add(u<<1|1,tr[u].lz3);
        tr[u].lz3=0;
    }
}

void up(int u)
{
    for(int t=0;t<=MAXK;t++)
        tr[u].s[t]=(tr[u<<1].s[t]+tr[u<<1|1].s[t])%P;
}

void build(int u,int l,int r)
{
    tr[u].len=r-l+1;
    tr[u].lz1=-1;
    tr[u].lz2=1;
    tr[u].lz3=0;
    if(l==r)
    {
        int pw=1;
        for(int t=0;t<=MAXK;t++)
        {
            tr[u].s[t]=pw;
            pw=1LL*pw*a[l]%P;
        }
        return;
    }
    int m=(l+r)>>1;
    build(u<<1,l,m);
    build(u<<1|1,m+1,r);
    up(u);
}

void upd_assign(int u,int l,int r,int ql,int qr,int x)
{
    if(ql<=l&&r<=qr) { push_assign(u,x);return; }
    pd(u);
    int m=(l+r)>>1;
    if(ql<=m)upd_assign(u<<1,l,m,ql,qr,x);
    if(qr>m)upd_assign(u<<1|1,m+1,r,ql,qr,x);
    up(u);
}

void upd_mul(int u,int l,int r,int ql,int qr,int c)
{
    if(ql<=l&&r<=qr) { push_mul(u,c);return; }
    pd(u);
    int m=(l+r)>>1;
    if(ql<=m)upd_mul(u<<1,l,m,ql,qr,c);
    if(qr>m)upd_mul(u<<1|1,m+1,r,ql,qr,c);
    up(u);
}

void upd_add(int u,int l,int r,int ql,int qr,int d)
{
    if(ql<=l&&r<=qr) { push_add(u,d);return; }
    pd(u);
    int m=(l+r)>>1;
    if(ql<=m)upd_add(u<<1,l,m,ql,qr,d);
    if(qr>m)upd_add(u<<1|1,m+1,r,ql,qr,d);
    up(u);
}

int qry(int u,int l,int r,int ql,int qr,int k)
{
    if(ql<=l&&r<=qr)return tr[u].s[k];
    pd(u);
    int m=(l+r)>>1,res=0;
    if(ql<=m)res=(res+qry(u<<1,l,m,ql,qr,k))%P;
    if(qr>m)res=(res+qry(u<<1|1,m+1,r,ql,qr,k))%P;
    return res;
}

signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    init();
    cin>>n>>m;
    for(int i=1;i<=n;i++)cin>>a[i];
    build(1,1,n);
    while(m--)
    {
        int op,l,r,x;
        cin>>op>>l>>r>>x;
        if(op==1)upd_assign(1,1,n,l,r,x);
        else if(op==2)upd_mul(1,1,n,l,r,x);
        else if(op==3)upd_add(1,1,n,l,r,x);
        else cout<<qry(1,1,n,l,r,x)<<'\n';
    }
    return 0;
}

NTT 优化

观察区间加的公式：

\[S_t' = \sum_{j=0}^{t} \binom{t}{j} d^{t-j} S_j \]

将其改写为指数生成函数形式：

\[\frac{S_t'}{t!} = \sum_{j=0}^{t} \frac{S_j}{j!} \cdot \frac{d^{t-j}}{(t-j)!} \]

令 \(A_j=\frac{S_j}{j!}\)，\(B_i=\frac{d^i}{i!}\)，则上式即为卷积形式：

\[S_t' = t! \cdot (A * B)_t \]

因此区间加法等价于计算序列 \(A\) 和 \(B\) 的卷积。使用 NTT 可在 \(O(K \log K)\) 时间内完成卷积计算。

模数 \(P = 998244353\)，原根为 \(G = 3\)。

时间复杂度

时间复杂度为 \(O(n\log n+qk\log k \log n)\)。

NTT 优化代码

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N=1e5+5,MAXK=100,P=998244353,G=3;
int n,m,a[N],C[MAXK+1][MAXK+1];
int rev[MAXK*4];
int fact[MAXK+1],invf[MAXK+1];
int qpow(int a,int b)
{
    int res=1;
    while(b)
    {
        if(b&1)res=1LL*res*a%P;
        a=1LL*a*a%P;
        b>>=1;
    }
    return res;
}
void init_ntt(int n)
{
    for(int i=0;i<n;i++)
        rev[i]=(rev[i>>1]>>1)|((i&1)?(n>>1):0);
}
void ntt(int *a,int n,int op)
{
    for(int i=0;i<n;i++)
        if(i<rev[i])swap(a[i],a[rev[i]]);

    for(int len=2;len<=n;len<<=1)
    {
        int wn=qpow(G,(P-1)/len);
        if(op==-1)wn=qpow(wn,P-2);
        for(int i=0;i<n;i+=len)
        {
            int w=1;
            for(int j=i;j<i+len/2;j++)
            {
                int u=a[j],v=1LL*w*a[j+len/2]%P;
                a[j]=(u+v)%P;
                a[j+len/2]=(u-v+P)%P;
                w=1LL*w*wn%P;
            }
        }
    }
    if(op==-1)
    {
        int inv=qpow(n,P-2);
        for(int i=0;i<n;i++)a[i]=1LL*a[i]*inv%P;
    }
}
struct Node
{
    int lz1,lz2,lz3,len,s[MAXK+1];
}tr[N<<2];
void init()
{
    for(int i=0;i<=MAXK;i++)
    {
        C[i][0]=C[i][i]=1;
        for(int j=1;j<i;j++)C[i][j]=(C[i-1][j-1]+C[i-1][j])%P;
    }
    fact[0]=1;
    for(int i=1;i<=MAXK;i++)fact[i]=1LL*fact[i-1]*i%P;
    invf[MAXK]=qpow(fact[MAXK],P-2);
    for(int i=MAXK-1;i>=0;i--)invf[i]=1LL*invf[i+1]*(i+1)%P;
}
void push_assign(int u,int x)
{
    tr[u].lz1=x;
    tr[u].lz2=1;
    tr[u].lz3=0;
    int pw=1;
    for(int t=0;t<=MAXK;t++)
    {
        tr[u].s[t]=1LL*tr[u].len*pw%P;
        pw=1LL*pw*x%P;
    }
}
void push_mul(int u,int c)
{
    tr[u].lz2=1LL*tr[u].lz2*c%P;
    tr[u].lz3=1LL*tr[u].lz3*c%P;
    int pw=1;
    for(int t=0;t<=MAXK;t++)
    {
        tr[u].s[t]=1LL*tr[u].s[t]*pw%P;
        pw=1LL*pw*c%P;
    }
}
void push_add(int u,int d)
{
    if(d==0)return;
    static int A[MAXK*4],B[MAXK*4];
    int k=MAXK;
    int n=1;
    while(n<=2*k)n<<=1;
    memset(A,0,n*sizeof(int));
    memset(B,0,n*sizeof(int));
    for(int j=0;j<=k;j++)
        A[j]=1LL*tr[u].s[j]*invf[j]%P;
    int pw=1;
    for(int i=0;i<=k;i++)
    {
        B[i]=1LL*pw*invf[i]%P;
        pw=1LL*pw*d%P;
    }
    init_ntt(n);
    ntt(A,n,1);
    ntt(B,n,1);
    for(int i=0;i<n;i++)A[i]=1LL*A[i]*B[i]%P;
    ntt(A,n,-1);
    for(int t=0;t<=k;t++)
        tr[u].s[t]=1LL*fact[t]*A[t]%P;
    tr[u].lz3=(tr[u].lz3+d)%P;
}
void pd(int u)
{
    if(tr[u].lz1!=-1)
    {
        push_assign(u<<1,tr[u].lz1);
        push_assign(u<<1|1,tr[u].lz1);
        tr[u].lz1=-1;
    }
    if(tr[u].lz2!=1)
    {
        push_mul(u<<1,tr[u].lz2);
        push_mul(u<<1|1,tr[u].lz2);
        tr[u].lz2=1;
    }
    if(tr[u].lz3)
    {
        push_add(u<<1,tr[u].lz3);
        push_add(u<<1|1,tr[u].lz3);
        tr[u].lz3=0;
    }
}
void up(int u)
{
    for(int t=0;t<=MAXK;t++)
        tr[u].s[t]=(tr[u<<1].s[t]+tr[u<<1|1].s[t])%P;
}
void build(int u,int l,int r)
{
    tr[u].len=r-l+1;
    tr[u].lz1=-1;
    tr[u].lz2=1;
    tr[u].lz3=0;
    if(l==r)
    {
        int pw=1;
        for(int t=0;t<=MAXK;t++)
        {
            tr[u].s[t]=pw;
            pw=1LL*pw*a[l]%P;
        }
        return;
    }
    int m=(l+r)>>1;
    build(u<<1,l,m);
    build(u<<1|1,m+1,r);
    up(u);
}
void upd_assign(int u,int l,int r,int ql,int qr,int x)
{
    if(ql<=l&&r<=qr) { push_assign(u,x);return; }
    pd(u);
    int m=(l+r)>>1;
    if(ql<=m)upd_assign(u<<1,l,m,ql,qr,x);
    if(qr>m)upd_assign(u<<1|1,m+1,r,ql,qr,x);
    up(u);
}
void upd_mul(int u,int l,int r,int ql,int qr,int c)
{
    if(ql<=l&&r<=qr) { push_mul(u,c);return; }
    pd(u);
    int m=(l+r)>>1;
    if(ql<=m)upd_mul(u<<1,l,m,ql,qr,c);
    if(qr>m)upd_mul(u<<1|1,m+1,r,ql,qr,c);
    up(u);
}
void upd_add(int u,int l,int r,int ql,int qr,int d)
{
    if(ql<=l&&r<=qr) { push_add(u,d);return; }
    pd(u);
    int m=(l+r)>>1;
    if(ql<=m)upd_add(u<<1,l,m,ql,qr,d);
    if(qr>m)upd_add(u<<1|1,m+1,r,ql,qr,d);
    up(u);
}
int qry(int u,int l,int r,int ql,int qr,int k)
{
    if(ql<=l&&r<=qr)return tr[u].s[k];
    pd(u);
    int m=(l+r)>>1,res=0;
    if(ql<=m)res=(res+qry(u<<1,l,m,ql,qr,k))%P;
    if(qr>m)res=(res+qry(u<<1|1,m+1,r,ql,qr,k))%P;
    return res;
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    init();
    cin>>n>>m;
    for(int i=1;i<=n;i++)cin>>a[i];
    build(1,1,n);
    while(m--)
    {
        int op,l,r,x;
        cin>>op>>l>>r>>x;
        if(op==1)upd_assign(1,1,n,l,r,x);
        else if(op==2)upd_mul(1,1,n,l,r,x);
        else if(op==3)upd_add(1,1,n,l,r,x);
        else cout<<qry(1,1,n,l,r,x)<<'\n';
    }
    return 0;
}

文章来源:https://www.cnblogs.com/BadBadBad/p/20899862/Segment-k
本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若内容造成侵权/违法违规/事实不符，请联系邮箱：jacktools123@163.com进行投诉反馈，一经查实，立即删除！

标签：

上一篇：AI Agent 30天速成｜Day7 教学笔记
下一篇：[开源] Vibe Coding 真心效率高， 100% Vibe Coding 的 WxPusher 客户端，支持 Mac， Window， Linux 。

相关文章

最新发布

点击排行

本站推荐

标签云