[BJOI2018] 求和 - 树上前缀和,LCA

Posted 2021-03-03 mollnn

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

Solution

对每个次数，预处理树上前缀和即可

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 300005;
const int mod = 998244353;

vector <int> g[N];
int n,m,t1,t2,t3,fa[N][20],s[N][55],dep[N],vis[N];

void dfs(int p) {
    vis[p]=1;
    for(int q:g[p]) if(vis[q]==0) {
        fa[q][0]=p;
        dep[q]=dep[p]+1;
        dfs(q);
    }
}

void dfs2(int p) {
    vis[p]=1;
    for(int q:g[p]) if(vis[q]==0) {
        for(int k=0;k<=50;k++) s[q][k]+=s[p][k], s[q][k]%=mod;
        dfs2(q);
    }
}

int lca(int p,int q) {
    if(dep[p]<dep[q]) swap(p,q);
    for(int i=18;i>=0;--i) if(dep[fa[p][i]]>=dep[q]) p=fa[p][i];
    for(int i=18;i>=0;--i) if(fa[p][i]-fa[q][i]) p=fa[p][i],q=fa[q][i];
    if(p-q) return fa[q][0];
    return p;
}

int dis(int p,int q,int k) {
    int l=lca(p,q);
    return ((s[p][k]+s[q][k]-s[l][k]-s[fa[l][0]][k])%mod+mod)%mod;
}

signed main() {
    ios::sync_with_stdio(false);
    cin>>n;
    for(int i=1;i<n;i++) {
        cin>>t1>>t2;
        g[t1].push_back(t2);
        g[t2].push_back(t1);
    }
    dfs(1);
    for(int i=1;i<=n;i++) {
        s[i][0]=1;
        for(int j=1;j<=50;j++) {
            s[i][j]=s[i][j-1]*dep[i]%mod;
        }
    }
    memset(vis,0,sizeof vis);
    dfs2(1);
    for(int i=1;i<=18;i++) {
        for(int j=1;j<=n;j++) {
            fa[j][i]=fa[fa[j][i-1]][i-1];
        }
    }
    cin>>m;
    for(int i=1;i<=m;i++) {
        cin>>t1>>t2>>t3;
        cout<<dis(t1,t2,t3)<<endl;
    }
}