树上前缀和

洛谷4427
思路
题目所求一段路径的幂次方之和,可以先进行预处理,再用相对应的前缀和进行求解。
首先先存储树,对树进行相对应的预处理,分清楚题目所求是点前缀和,还是边前缀和,两者的公式有所差异。
点前缀和是 一端前缀和+另一端前缀和-(lca的前缀和+lca的父节点的前缀和)。
边前缀和是 一端前缀和+另一端前缀和-2*lca的前缀和。
思路来自董晓算法

点击查看代码
#include<bits/stdc++.h>
using namespace std;
 #define int long long
#define ll long long
vector<int>v[300005];
int n,k;
int fa[300005][21];
int m=998244353;
int p[51],sum[300005][51];
int dev[300005]={0};
void dfs(int u,int father)
{

    for(int i=1;i<=20;i++)
    {
     //if(fa[u][i-1]!=0)
      fa[u][i]=fa[fa[u][i-1]][i-1];
    }
    for(int i:v[u])
    {
        if(i==father)
            continue;
     dev[i]=dev[u]+1;
     fa[i][0]=u;
     for(int j=1;j<=50;j++)
     {
      p[j]=(p[j-1]%m*dev[i])%m;
 
     }
     for(int j=1;j<=50;j++)
     {
       sum[i][j]=sum[u][j]+p[j]%m;
     }
        dfs(i,u);
    }

}
int lca(int x,int y)
{
    if(dev[x]<dev[y])
    swap(x,y);
    for(int i=20;i>=0;i--)
    {
        if(dev[fa[x][i]]>=dev[y])
            x=fa[x][i];
    }
    if(x==y)
        return y;
    for(int i=20;i>=0;i--)
    {
        if(fa[x][i]!=fa[y][i])
        {
            x=fa[x][i];
            y=fa[y][i];

        }

    }
     return fa[x][0];

}

signed main()
{
     cin.tie(NULL);
     cout.tie(nullptr);
     ios_base::sync_with_stdio(false);
     cin>>n;
     for(int i=1;i<=n-1;i++)
     {
        int x,y;
        cin>>x>>y;
        v[x].push_back(y);
        v[y].push_back(x);
     }
     cin>>k;
     p[0]=1;
     dfs(1,0);
     //cout<<lca(4,5)<<" j jj"<<'\n';
     for(int i=1;i<=k;i++)
     {
         int s,t,z;
         cin>>s>>t>>z;
         int l=lca(s,t);
         cout<<(sum[s][z]+sum[t][z]-sum[l][z]-sum[fa[l][0]][z]+m)%m<<'\n';
     }


   return 0;
}

posted @ 2024-01-30 00:25  WandW  阅读(80)  评论(0)    收藏  举报