bzoj5293: [Bjoi2018]求和(Lca)

Hanks_o發表於2018-04-18

題目傳送門

解法:
k<=50
暴力算每個點到根的每個次方和。
然後求出lca差分減一下

程式碼實現:

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#include<cmath>
#include<queue>
using namespace std;
typedef long long ll;
struct node {int x,y,next;}a[610000];int len,last[310000];
void ins(int x,int y) {len++;a[len].x=x;a[len].y=y;a[len].next=last[x];last[x]=len;}
int mx[21][310000],dep[310000],n,D;
void dfs(int x) {
    D=max(D,dep[x]);
    for(int k=last[x];k;k=a[k].next) {
        int y=a[k].y;
        if(y!=mx[0][x]) {mx[0][y]=x;dep[y]=dep[x]+1;dfs(y);}
    }
}
int bin[21];
void work() {
    bin[0]=1;for(int i=1;i<=20;i++)bin[i]=bin[i-1]*2;
    for(int j=1;j<=20;j++)for(int i=1;i<=n;i++)if(dep[i]>=bin[j]) 
        mx[j][i]=mx[j-1][mx[j-1][i]];
}
int lca(int x,int y) {
    if(dep[x]>dep[y])swap(x,y);
    for(int i=20;i>=0;i--)if(dep[y]-dep[x]>=bin[i])y=mx[i][y];
    if(x==y)return x;
    for(int i=20;i>=0;i--)if(dep[x]>=bin[i]&&mx[i][x]!=mx[i][y])
        x=mx[i][x],y=mx[i][y];
    return mx[0][x];
}
const ll mod=ll(998244353);
ll p[310000][52],sum[310000][52];
void merge(int x,int c) {
    for(int k=last[x];k;k=a[k].next) {
        int y=a[k].y;
        if(y!=mx[0][x]) {sum[y][c]=(sum[x][c]+p[dep[y]][c])%mod;merge(y,c);}
    }
}
int main() {
    scanf("%d",&n);len=0;memset(last,0,sizeof(last));
    for(int i=1;i<n;i++) {int x,y;scanf("%d%d",&x,&y);ins(x,y);ins(y,x);}
    D=0;dep[1]=0;mx[0][1]=0;dfs(1);work();
    for(int i=1;i<=D;i++) {
        p[i][1]=(ll)i;
        for(int j=2;j<=50;j++) p[i][j]=ll(p[i][j-1]*i)%mod;
    }memset(sum,0,sizeof(sum));
    for(int i=1;i<=50;i++)merge(1,i);
    int Q;scanf("%d",&Q);
    while(Q--) {
        int x,y,k;scanf("%d%d%d",&x,&y,&k);
        int Lca=lca(x,y);
        printf("%lld\n",((sum[x][k]+sum[y][k]-sum[Lca][k]-sum[mx[0][Lca]][k])%mod+mod)%mod);
    }
    return 0;
}