$Luogu$ $P4427$ $[BJOI2018]$ 求和

链接

背景

\(CCF\) \(NOI\) \(2018\) 北京市代表队选拔 \(Day1\) \(T1\) ， \(Luogu\) \(P4427/LOJ2491\)

题意

给定 \(n\) 个点构成的一棵以 \(1\) 号点为根的树，定义根节点的深度为 \(0\) ， \(m\) 次询问 \(x,y\) 两点间的所有点的深度的 \(k\) 次方和。

解法

由于 \(k \leqslant 50\) ，不难想到 \(bfs\) 预处理好每个点 \(x\) 的深度 \(d_x\) 以及从根节点开始到 \(x\) 点为止的权值的 \(k\) 次方和 \(dis_{x,k}\) （不包含根节点的权值）。
那么，相当于每个点定义了一个新的距离，问题又变成了树上两点间的距离。倍增法求最近公共祖先即可。
设 \(x,y\) 两点的最近公共祖先是 \(pos\) ，则答案是 \(dis_{x,k}+dis_{y,k}-dis_{pos,k}-dis{f_{pos,0},k}\) 。因为加上的是两遍根节点除外到 \(pos\) 点路上所有点的新权值和以及 \(pos\) 点除外到 \(x\) 点和 \(y\) 点路上所有点的新权值和，减去的是一遍根节点除外到 \(pos\) 点路上所有点的新权值和以及一遍根节点除外到 \(pos\) 点的父节点路上所有点的新权值和，所以不重不漏。（容斥原理）
此题本质上就是在倍增法求最近公共祖先模板上换成了点权、重新定义了距离。

细节

（划重点）请务必务必务必注意在倍增的题目写完前拿最大的数据检查倍增数组小的一维开的够不够大！！！

代码

\(View\) \(Code\)

#include<bits/stdc++.h>
using namespace std;
inline int read()
{
    int ret=0,f=1;
    char ch=getchar();
    while(ch>'9'||ch<'0')
    {
        if(ch=='-')
            f=-1;
        ch=getchar();
    }
    while(ch>='0'&&ch<='9')
    {
        ret=(ret<<1)+(ret<<3)+ch-'0';
        ch=getchar();
    }
    return ret*f;
}
const int mod=998244353;
int n,m,t,x,y,k;
int f[300005][20],d[300005];
int num,head[600005];
long long dis[300005][55];
queue<int> q;
struct edge
{
    int ver,nxt;
}e[600005];
inline void adde(int u,int v)
{
    e[++num].ver=v;
    e[num].nxt=head[u];
    head[u]=num;
}
inline long long qpow(long long a,long long b)
{
    long long ans=1;
    while(b)
    {
        if(b&1)
            ans=ans*a%mod;
        a=a*a%mod;
        b>>=1;
    }
    return ans%mod;
}
inline void bfs()
{
    d[1]=1;
    q.push(1);
    while(!q.empty())
    {
        int x=q.front();
        q.pop();
        for(register int i=head[x];i;i=e[i].nxt)
        {
            int y=e[i].ver;
            if(d[y])
                continue;
            d[y]=d[x]+1;
            for(register int j=1;j<=50;j++)
                dis[y][j]=(dis[x][j]+qpow(d[x],j))%mod;
            f[y][0]=x;
            for(register int j=1;j<=t;j++)
                f[y][j]=f[f[y][j-1]][j-1];
            q.push(y);
        }
    }
}
inline int lca(int x,int y)
{
    if(d[x]<d[y])
        swap(x,y);
    for(register int i=t;i||!i;i--)
        if(d[y]<=d[f[x][i]])
            x=f[x][i];
    if(x==y)
        return y;
    for(register int i=t;i||!i;i--)
    {
        if(f[x][i]!=f[y][i])
        {
            x=f[x][i];
            y=f[y][i];
        }
    }
    return f[x][0];
}
int main()
{
    n=read();
    t=(int)(log(n)/log(2))+1;
    for(register int i=1;i<n;i++)
    {
        x=read();
        y=read();
        adde(x,y);
        adde(y,x);
    }
    bfs();
    m=read();
    for(register int i=1;i<=m;i++)
    {
        x=read();
        y=read();
        k=read();
        int pos=lca(x,y);
        long long ans=((dis[x][k]+dis[y][k]-dis[pos][k]-dis[f[pos][0]][k])%mod+mod)%mod;
        printf("%lld\n",ans);
    }
    return 0;
}