BZOJ 5293 求和(LCA)

5293: [Bjoi2018]求和

Time Limit: 20 Sec  Memory Limit: 512 MB
Submit: 203  Solved: 128
[Submit][Status][Discuss]

Description

master 对树上的求和非常感兴趣。他生成了一棵有根树，并且希望多次询问这棵树上一段路径上所有节点深度的k

 次方和，而且每次的k 可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给

了pupil，但pupil 并不会这么复杂的操作，你能帮他解决吗？

Input

第一行包含一个正整数n ，表示树的节点数。

之后n-1 行每行两个空格隔开的正整数i,j ，表示树上的一条连接点i 和点j 的边。

之后一行一个正整数m ，表示询问的数量。

之后每行三个空格隔开的正整数i,j,k ，表示询问从点i 到点j 的路径上所有节点深度的k 次方和。

由于这个结果可能非常大，输出其对998244353 取模的结果。

树的节点从1 开始标号，其中1 号节点为树的根。

Output

对于每组数据输出一行一个正整数表示取模后的结果。

1≤n,m≤300000,1≤k≤50

Sample Input

5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45

Sample Output

33
503245989

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<cmath>
#include <queue>
#include <iostream>

using namespace std;
typedef long long LL;
#define N 300005
const int mod = 998244353;
int pow(int a,int b){
    int odd = 1;
    while (b){
        if(b&1){
            odd = (int)(1ll*odd*a%mod);
        }
        a = (int)(1ll*a*a%mod);
        b>>=1;
    }
    return odd;
}
class Graphic {
public:
    Graphic(){
        bin[0] =1;
        for(int i = 1 ; i < 25 ; ++i)bin[i] = bin[i-1]<<1;
    }
    void init(){
        cnt = 0;
        dep[1] = 1;
        fa[1][0] = 0;
        memset(head,-1, sizeof(head));
    }
    void pre_deal(){
        dfs(1);
        cal(1);
    }
    void add(int u,int v){
        _add(u,v);
        _add(v,u);
    }
    int find_sum(int x,int y,int rat){
        int la = lca(x,y);
        return int((nx[x][rat]+nx[y][rat]-nx[la][rat]-nx[fa[la][0]][rat])%mod);
    }
private:
    struct Edge {
        int to, next;
    } edge[N << 1];
    int head[N],cnt;
    int bin[25];
    int dep[N],
        fa[N][25];///fa[i][j]表示 i的第 (1<<j)个父亲
    LL nx[N][52];///nx[i][j] 表示节点i的j次方和
    void _add(int u,int v){
        edge[cnt].to = v;
        edge[cnt].next = head[u];
        head[u] = cnt++;
    }
    int lca(int x,int y){
        if(dep[x]<dep[y])swap(x,y);
        for(int i=20;i>=0;i--)if(bin[i]<=dep[x] and dep[fa[x][i]]>=dep[y])x=fa[x][i];///将x y的深度调成相同的
        if(x==y)return x;
        for(int i=20;i>=0;i--)if(bin[i]<=dep[x] and fa[x][i]!=fa[y][i])x=fa[x][i],y=fa[y][i];
        return fa[x][0];
    }
    void cal(int x){
        for(int i = 1 ; i <= 50 ; ++i)nx[x][i] = nx[fa[x][0]][i]+pow(dep[x]-1,i);
        int v;
        for(int i = head[x] ; ~i ; i = edge[i].next){
            v = edge[i].to;
            if(v!=fa[x][0])cal(v);
        }
    }
    void dfs(int x){
        for(int i = 1 ; i <= 20 ; ++i){
            if(fa[x][i-1]){
                fa[x][i] = fa[fa[x][i-1]][i-1];
                /// x的第 2^i个父亲就是 x的2^(i-1)个父亲的2^(i-1)个父亲
            }
            else break;
        }
        int v;
        for(int i = head[x] ; ~i ;i = edge[i].next ){
            v = edge[i].to;
            if(v!=fa[x][0]){
                fa[v][0] = x;
                dep[v] = dep[x]+1;
                dfs(v);
            }
        }
    }
};
Graphic ss;
int main(){
    int n,m,a,b;
    while(cin>>n){
        ss.init();
        for(int i = 1 ; i < n ; ++i){
            scanf("%d %d",&a,&b);
            ss.add(a,b);
        }
        ss.pre_deal();
        cin>>m;
        int k;
        while (m--){
            scanf("%d %d %d",&a,&b,&k);
            printf("%d
",ss.find_sum(a,b,k));
        }
    }
}