題目描述:
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);
}
}