UPC-6744 求和VII(LCA倍增法+幂次前缀和)

题目描述
master对树上的求和非常感兴趣。他生成了一棵有根树,并且希望多次询问这棵树上一段路径上所有节点深度的k次方和,而且每次的k可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给了pupil,但pupil并不会这么复杂的操作,你能帮他解决吗?

输入
第一行包含一个正整数n,表示树的节点数。
之后n−1行每行两个空格隔开的正整数i,j,表示树上的一条连接点i和点j的边。
之后一行一个正整数m,表示询问的数量。
之后每行三个空格隔开的正整数i,j,k,表示询问从点i到点j的路径上所有节点深度的k次方和。由于这个结果可能非常大,输出其对998244353取模的结果。
树的节点从1开始标号,其中1号节点为树的根。

输出
对于每组数据输出一行一个正整数表示取模后的结果。

样例输入
5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45

样例输出
33
503245989

提示
以下用d(i)表示第i个节点的深度。
对于样例中的树,有d(1)=0,d(2)=1,d(3)=1,d(4)=2,d(5)=2。
因此第一个询问答案为(25+15+05) mod 998244353=33,第二个询问答案为(245+145+245) mod 998244353=503245989。

对于30%的数据,1≤n,m≤100;
对于60%的数据,1≤n,m≤1000;
对于100%的数据,1≤n,m≤300000,1≤k≤50。

思路:首先很裸的题意就是要找到两结点路径上的所有深度幂次和,那么即要求LCA,找到路径两端点上的深度,找到LCA的深度,然后直接计算经过深度的幂次求和即可。

这里的询问次数很多,因此不能暴力的用朴素LCA,倍增法找到最近公共祖先后也不能用for循环求和经过的深度。明显会超时。因此,我们注意到幂次范围很小,只有50,并且确定了树结构后,我们已知树的最深深度。那么预处理好所有深度的1~50幂次结果,并且处理成前缀和的形式存储,这样在知道LCA和端点的情况下就可以直接作差以O(1)的复杂度求出一半路径的深度幂次和。另一半同理也是作差,或者,直接用一个端点的前缀和加另一个端点的前缀和,此时发现从根节点到LCA的这段重复加了2次,那么减去,LCA的前缀和,这时LCA本身被减去了,因此加回来一次,这样求和即路径上所有节点的的深度幂次和。

写这题时LCA倍增模板抄错了找了几乎一个小时的BUG。。。。

#include<bits/stdc++.h>
#define LL long long
#define M(a,b) memset(a,b,sizeof a)
#define pb(x) push_back(x)
using namespace std;
const int maxn=3e5+7;
const int mod=998244353;
int n,t,mxdep=0;
LL k;
int fa[23][maxn],dep[maxn];
LL sum[maxn][55];
vector<int>g[maxn];
LL qp(LL a,LL b)
{
    if(a==0)return 0;
    LL ans=1;
    while(b)
    {
        if(b&1)ans=(ans*a)%mod;
        a=(a*a)%mod;
        b>>=1;
    }
    return ans;

}
void init()
{
    for(int i=0; i<20; i++)
        for(int j=1; j<=n; j++)
        {
            if(fa[i][j]==-1)continue;
            fa[i+1][j]=fa[i][fa[i][j]];
        }

    for(LL i=1; i<=50; i++)///预处理所有深度幂次为1~50的值的前缀和
        for(LL j=0; j<=mxdep; j++)
            if(j==0)sum[j][i]=0;
            else sum[j][i]=(sum[j-1][i]+qp(j,i))%mod;
}
void dfs(int u)///深搜计算深度
{
    if(g[u].size()==0)return;
    for(int i=0; i<g[u].size(); i++)
    {
        int v=g[u][i];
        dep[v]=dep[u]+1;
        if(dep[v]>mxdep)mxdep=dep[v];
        dfs(v);
    }
}
int lca(int u,int v)///倍增找最近公共祖先
{
    if(dep[u]>dep[v]) swap(u,v);
    for(int i=20; i>=0; i--)
        if((dep[v]-dep[u])&(1<<i))
            v=fa[i][v];
    if(u==v)return u;
    for(int i=20; i>=0; i--)
        if(fa[i][u]!=fa[i][v])
        {
            u=fa[i][u];
            v=fa[i][v];
        }
    return fa[0][u];
}
int main()
{
    mxdep=0;
    scanf("%d",&n);
    for(int i=0; i<=n; i++) g[i].clear();
    M(fa,-1);
    M(dep,-1);
    int u,v;
    for(int i=1; i<n; i++)
    {
        scanf("%d%d",&u,&v);
        fa[0][v]=u;
        g[u].pb(v);
    }
    int rt;
    for(int i=1; i<=n; i++)
        if(fa[0][i]==-1)
        {
            rt=i;
            break;
        }
    dep[rt]=0;
    dfs(rt);
    init();
    scanf("%d",&t);
    while(t--)
    {
        scanf("%d%d%lld",&u,&v,&k);
        int tmp=lca(u,v);///拿路径两端点的深度与LCA的深度求和
        LL  ans=(sum[dep[u]][k]+sum[dep[v]][k])%mod;///利用前缀和快速求出两深度上经过的所有深度幂次和
        ans=(ans-sum[dep[tmp]][k]+mod)%mod;///减去两次公共祖先的前缀和,加上一次公共祖先的幂次
        ans=(ans-sum[dep[tmp]][k]+mod)%mod;
        ans=(ans+qp(dep[tmp],k))%mod;
        printf("%lld
",ans);
    }
}
/*
5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45
*/
原文地址:https://www.cnblogs.com/kuronekonano/p/11135722.html