A09【模板】树上前缀和 P4427 [BJOI2018] 求和

A09【模板】树上前缀和 P4427 [BJOI2018] 求和_哔哩哔哩_bilibili

 

前缀和 & 差分 - OI Wiki

 

D09 倍增算法 P3379【模板】最近公共祖先（LCA） - 董晓 - 博客园

P4427 [BJOI2018] 求和 - 洛谷

// 树上前缀和 O(nlogn)
#include<bits/stdc++.h>
using namespace std;

typedef long long LL;
const LL mod=998244353;
const int N=3e5+10;
int n,m;
int h[N],to[2*N],ne[2*N],tot;
void add(int a,int b){ 
  to[++tot]=b,ne[tot]=h[a],h[a]=tot;
}

int fa[N][22]; //fa[x][i]表示从x向上跳2^i层的祖先结点
int dep[N];    //dep[y]表示y的深度
LL mi[55];     //mi[j]表示dep[y]的j次幂
LL s[N][55];   //s[y][j]表示从根到y的路径节点的深度的j次幂之和

void dfs(int x,int f){ 
  for(int i=1; i<=20; i++) fa[x][i]=fa[fa[x][i-1]][i-1];
  for(int i=h[x];i;i=ne[i]){
    int y=to[i];
    if(y!=f){
      fa[y][0]=x; dep[y]=dep[x]+1;
      for(int j=1;j<=50;j++) mi[j]=mi[j-1]*dep[y]%mod;
      for(int j=1;j<=50;j++) s[y][j]=(mi[j]+s[x][j])%mod;
      dfs(y,x);
    }
  }
}
int lca(int x,int y){
  if(dep[x]<dep[y])swap(x,y);
  for(int i=20; ~i; i--)if(dep[fa[x][i]]>=dep[y]) x=fa[x][i];
  if(x==y) return y;
  for(int i=20; ~i; i--)if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
  return fa[x][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); //倍增预处理 dep,fa,mi,s 数组
  scanf("%d",&m);
  for(int i=1,x,y,k;i<=m;i++){
    scanf("%d%d%d",&x,&y,&k);
    int z=lca(x,y); //倍增求lca
    LL ans=(s[x][k]+s[y][k]-s[z][k]-s[fa[z][0]][k]+2*mod)%mod;
    printf("%lld\n",ans);
  }
}