[BJOI2018]求和

考虑每一次要求的是一条路径上深度的\(k\)次方和,但是\(k \leq 50\),所以可以先把\(k\)的所有情况预处理出来,然后直接树上差分就行了。

/*
  mail: mleautomaton@foxmail.com
  author: MLEAutoMaton
  This Code is made by MLEAutoMaton
*/
#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<math.h>
#include<algorithm>
#include<queue>
#include<set>
#include<map>
#include<iostream>
using namespace std;
#define ll long long
#define re register
#define file(a) freopen(a".in","r",stdin);freopen(a".out","w",stdout)
inline int gi(){
    int f=1,sum=0;char ch=getchar();
    while(ch>'9' || ch<'0'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0' && ch<='9'){sum=(sum<<3)+(sum<<1)+ch-'0';ch=getchar();}
    return f*sum;
}
const int N=300010,Mod=998244353;
int n,front[N],cnt,f[N][23],dep[N],dis[N][60];
struct node{int to,nxt;}e[N<<1];
int qpow(int a,int b){int ret=1;while(b){if(b&1)ret=1ll*ret*a%Mod;b>>=1;a=1ll*a*a%Mod;}return ret;}
void Add(int u,int v){e[++cnt]=(node){v,front[u]};front[u]=cnt;}
void dfs(int u,int ff){
    dep[u]=dep[ff]+1;f[u][0]=ff;
    for(int k=1;k<=50;k++)dis[u][k]=(dis[ff][k]+qpow(dep[u],k))%Mod;
    for(int i=front[u];i;i=e[i].nxt){
        int v=e[i].to;if(v==ff)continue;
        dfs(v,u);
    }
}
int getlca(int x,int y){
    if(dep[x]<dep[y])swap(x,y);
    for(int i=20;~i;i--)
        if(dep[x]-(1<<i)>=dep[y])x=f[x][i];
    if(x==y)return x;
    for(int i=20;~i;i--)
        if(f[x][i]!=f[y][i])x=f[x][i],y=f[y][i];
    return f[x][0];
}
int solve(int x,int y,int k){
    int lca=getlca(x,y);
    return (((dis[x][k]+dis[y][k])%Mod-dis[lca][k]+Mod)%Mod-dis[f[lca][0]][k]+Mod)%Mod;
}
int main(){
    n=gi();
    for(int i=1;i<n;i++){
        int u=gi(),v=gi();
        Add(u,v);Add(v,u);
    }
    dep[1]=-1;dfs(1,1);
    f[1][0]=0;
    for(int j=1;j<=20;j++)
        for(int i=1;i<=n;i++)
            f[i][j]=f[f[i][j-1]][j-1];
    int m=gi();
    while(m--){
        int u=gi(),v=gi(),k=gi();
        printf("%d\n",solve(u,v,k));
    }
    return 0;
}
上一篇:[洛谷 P4457] [BJOI2018]治疗之雨


下一篇:[BZOJ5294][BJOI2018]二进制(线段树)