P4427 [BJOI2018]求和（倍增LCA、树链剖分）

题目描述

master 对树上的求和非常感兴趣。他生成了一棵有根树，并且希望多次询问这棵树上一段路径上所有节点深度的kk 次方和，而且每次的kk 可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给了pupil，但pupil 并不会这么复杂的操作，你能帮他解决吗？

输入格式

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

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

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

之后每行三个空格隔开的正整数i, j, ki,j,k，表示询问从点ii 到点jj 的路径上所有节点深度的kk 次方和。由于这个结果可能非常大，输出其对998244353998244353 取模的结果。

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

输出格式

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

题解：

注意到k的范围只有50，所以先预处理出每个节点的深度k次前缀和，然后用树上差分的思想求出答案。求LCA用树剖或倍增LCA都可以。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=3e5+100;
const ll mod=998244353;
ll h[55][maxn];
int n,m;
vector<int> g[maxn];

int father[20][maxn];
ll db[55][maxn];//预处理所有数的k次
 
void dfs (int x) {
    for (int i=0;i<g[x].size();i++) {
        int v=g[x][i];
        if (v==father[0][x]) continue;
        father[0][v]=x;
        h[0][v]=h[0][x]+1;
        for (int k=1;k<=50;k++)
            h[k][v]=h[k][x]+db[k][h[0][v]],
            h[k][v]%=mod;
        dfs(v);
    }
}
int lca (int x,int y) {
    if (h[0][x]<h[0][y]) swap(x,y);
    for (int i=17;i>=0;i--)
        if (h[0][x]-h[0][y]>>i) x=father[i][x];
    if (x==y) return x;
    for (int i=17;i>=0;i--) 
        if (father[i][x]!=father[i][y]) {
            x=father[i][x];
            y=father[i][y];
        }
    return father[0][x];
}
ll getDis (int x,int y,int k) {
    int Lca=lca(x,y);
    int lf=father[0][Lca];
    return (h[k][x]+h[k][y])%mod+2*mod-(h[k][Lca]+h[k][lf]);
}

int main () {
    scanf("%d",&n);
    for (int i=1;i<=n;i++) {
        db[1][i]=i;
        for (int j=2;j<=50;j++)    
            db[j][i]=db[j-1][i]*db[1][i],
            db[j][i]%=mod;
    }
    for (int i=1;i<n;i++) {
        int x,y;
        scanf("%d%d",&x,&y);
        g[x].push_back(y);
        g[y].push_back(x);
    }
    dfs(1);
    for (int i=1;i<=17;i++)
        for (int j=1;j<=n;j++)
            father[i][j]=father[i-1][father[i-1][j]];
    scanf("%d",&m);
    for (int i=1;i<=m;i++) {
        int x,y,k;
        scanf("%d%d%d",&x,&y,&k);
        printf("%lld\n",(getDis(x,y,k)+mod)%mod);
    }
}