题目描述:
master 对树上的求和非常感兴趣。他生成了一棵有根树,并且希望多次询问这棵树上一段路径上所有节点深度的k
次方和,而且每次的k 可能是不同的。此处节点深度的定义是这个节点到根的路径上的边数。他把这个问题交给
了pupil,但pupil 并不会这么复杂的操作,你能帮他解决吗?
输入:
第一行包含一个正整数n ,表示树的节点数。
之后n-1 行每行两个空格隔开的正整数i,j ,表示树上的一条连接点i 和点j 的边。
之后一行一个正整数m ,表示询问的数量。
之后每行三个空格隔开的正整数i,j,k ,表示询问从点i 到点j 的路径上所有节点深度的k 次方和。
由于这个结果可能非常大,输出其对998244353 取模的结果。
树的节点从1 开始标号,其中1 号节点为树的根。
输出:
对于每组数据输出一行一个正整数表示取模后的结果。
1≤n,m≤300000,1≤k≤50
样例输入:
5
1 2
1 3
2 4
2 5
2
1 4 5
5 4 45
样例输出:
33
503245989
说明:
样例解释 以下用d(i) 表示第i 个节点的深度。 对于样例中的树,有d(1)=0,d(2)=1,d(3)=1,d(4)=2,d(5)=2。
因此第一个询问答案为(2^5 + 1^5 + 0^5) mod 998244353 = 33 第二个询问答案为(2^45 + 1^45 + 2^45) mod 998244353 = 503245989。
思路分析:
我们可以用一个sum[i][j]数组来存储节点i到根节点的路径的j次。
而这怎么求?
我们可以在处理树时,用它父节点的数组数来处理就可以了(详情见代码);
最后我们可以用(((sum[a][c]+sum[b][c])%mod-sum[lc][c]+mod)%mod-sum[fa[lc][0]][c]+mod)%mod求出答案。
代码实现:
#include<cstdio>
#include<iostream>
#include<vector>
#include<cstring>
using namespace std;
#define ll long long
#define mod 998244353
int n,m;
ll sum[300005][55];
int dep[300005],fa[300005][25];
vector<int>G[300005];
void dfs(int u,int Fa)
{
if (u != 1)
dep[u]=dep[Fa]+1;
fa[u][0] = Fa;
ll ans = 1;
for(int i=1;i<=50;i++){
ans = ans * dep[u] % mod;
sum[u][i]=(ans+sum[fa[u][0]][i])%mod;
}
for(int i=0;i<=20;i++)
fa[u][i+1]=fa[fa[u][i]][i];
for(int i=0;i<G[u].size();i++)
{
if(Fa==G[u][i])
continue;
dfs(G[u][i],u);
}
}
int lca(int u,int v)
{
if(dep[u]<dep[v])
swap(u,v);
for(int i=20;i>=0;i--)
{
if(dep[fa[u][i]]>=dep[v])
u=fa[u][i];
if(u==v)
return u;
}
for(int i=20;i>=0;i--)
{
if(fa[u][i]!=fa[v][i])
{
u=fa[u][i];
v=fa[v][i];
}
}
return fa[u][0];
}
int main()
{
scanf("%d",&n);
for(int i=1;i<n;i++)
{
int a,b;
scanf("%d%d",&a,&b);
G[a].push_back(b);
G[b].push_back(a);
}
scanf("%d",&m);
dfs(1,0);
for(int i=1;i<=m;i++)
{
int a,b,c;
scanf("%d%d%d",&a,&b,&c);
int lc=lca(a,b);
ll ans=(((sum[a][c]+sum[b][c])%mod-sum[lc][c]+mod)%mod-sum[fa[lc][0]][c]+mod)%mod;
printf("%lld\n",ans);
}
}