[BZOJ5293][BJOI2018]求和(倍增)

裸的树上倍增。

#include<cstdio>
#include<cstring>
#include<algorithm>
#define rep(i,l,r) for (int i=(l); i<=(r); i++)
using namespace std;

const int N=300010,mod=998244353;
struct E{ int to,nxt; }e[N*2];
int fa[N],a[N],dep[N],f[N][21],sm[N][52],pow[N][52];
int h[N],cnt,n,m,x,y,k;

void add(int u,int v){ e[++cnt].to=v; e[cnt].nxt=h[u]; h[u]=cnt; }

void Build(int x){
    dep[x]=dep[fa[x]]+1;
    sm[x][0]=1; pow[x][0]=1;
    rep(i,1,50){
        pow[x][i]=((long long)pow[x][i-1]*(long long)dep[x])%mod;
        sm[x][i]=(sm[fa[x]][i]+pow[x][i])%mod;
    }
    for (int i=h[x]; i; i=e[i].nxt)
        if (e[i].to!=fa[x]){
            fa[e[i].to]=x;
            f[e[i].to][0]=x;
            Build(e[i].to);
        }
}

int LCA(int x,int y){
    if (dep[x]<dep[y]) swap(x,y);
    for (int i=19; i>=0; --i)
        if (dep[f[x][i]]>=dep[y]) x=f[x][i];
    if (x==y) return x;
    for (int i=19; i>=0; --i)
        if (f[x][i]!=f[y][i]) x=f[x][i],y=f[y][i];
    return fa[x];
}

int main(){
    freopen("sum.in","r",stdin);
    freopen("sum.out","w",stdout);
    scanf("%d",&n);
    rep(i,1,n-1) scanf("%d%d",&x,&y),add(x,y),add(y,x);
    dep[0]=-1; Build(1);
    rep(i,1,19) rep(j,1,n) f[j][i]=f[f[j][i-1]][i-1];
    scanf("%d",&m);
    rep(i,1,m){
        scanf("%d%d%d",&x,&y,&k); int lca=LCA(x,y);
        printf("%d
",((sm[x][k]+sm[y][k]-sm[fa[lca]][k]-sm[lca][k])%mod+mod)%mod);
    }
    return 0;
}