A09 树上前缀和 P4427 [BJOI2018] 求和

视频链接：A09 树上前缀和 P4427 [BJOI2018] 求和_哔哩哔哩_bilibili

Luogu P4427 [BJOI2018] 求和

 

 

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;

typedef long long LL;
const LL mod=998244353;
const int N=3e5+10,M=2*N; 
int n,m;
int h[N],to[M],ne[M],tot;
int fa[N][22];//fa[u][i]表示从u向上跳2^i层的祖先结点
int dep[N];   //dep[v]表示v的深度
LL mi[60];   //mi[j]表示dep[v]的j次幂
LL s[N][60]; //s[v][j]表示从根到v的路径节点的深度的j次幂之和

void add(int a, int b){ 
  to[++tot]=b,ne[tot]=h[a],h[a]=tot;
}
void dfs(int u, int f){ //预处理节点信息
  for(int i=1; i<=20; i++) 
    fa[u][i]=fa[fa[u][i-1]][i-1];
  for(int i=h[u];i;i=ne[i]){
    int v=to[i];
    if(v==f) continue;
    fa[v][0]=u; dep[v]=dep[u]+1;
    for(int j=1;j<=50;j++) mi[j]=mi[j-1]*dep[v]%mod;
    for(int j=1;j<=50;j++) s[v][j]=(mi[j]+s[u][j])%mod;
    dfs(v, u);
  }
}
int lca(int u, int v){ //倍增法求lca
  if(dep[u]<dep[v])swap(u, v);
  for(int i=20; ~i; i--)
    if(dep[fa[u][i]]>=dep[v]) u=fa[u][i];
  if(u==v) return v;
  for(int i=20; ~i; i--)
    if(fa[u][i]!=fa[v][i]) u=fa[u][i], v=fa[v][i];
  return fa[u][0];
}
int main(){
  scanf("%d",&n); 
  for(int i=1,a,b; i<n; i++){
    scanf("%d%d",&a,&b);
    add(a,b); add(b,a);
  }
  mi[0]=1; dfs(1, 0);
  
  scanf("%d",&m);
  for(int i=1,u,v,k;i<=m;i++){
    scanf("%d%d%d",&u,&v,&k);
    int l=lca(u,v);
    LL ans=(s[u][k]+s[v][k]-s[l][k]-s[fa[l][0]][k]+2*mod)%mod;
    printf("%lld\n",ans);
  }
}