C++最近公共祖先(LCA)例題—————求和

題目描述:

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);
	}
}

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章