金題大戰Vol.0 C、樹上的等差數列

題目描述

給定一棵包含\(N\)個節點的無根樹，節點編號\(1-N\)。其中每個節點都具有一個權值，第\(i\)個節點的權值是\(A_i\)。

小\(Hi\)希望你能找到樹上的一條最長路徑，滿足沿著路徑經過的節點的權值序列恰好構成等差數列。

輸入格式

第一行包含一個整數\(N\)。

第二行包含\(N\)個整數\(A_1, A_2, ... A_N\)。

以下\(N-1\)行，每行包含兩個整數\(U\)和\(V\)，代表節點\(U\)和\(V\)之間有一條邊相連。

輸出格式

最長等差數列路徑的長度

樣例

樣例輸入

7
3 2 4 5 6 7 5
1 2
1 3
2 7
3 4
3 5
3 6

樣例輸出

4

資料範圍與提示

對於\(50\%\)的資料，\(1 ≤ N ≤ 1000\)

對於\(100\%\)的資料，\(1 ≤ N ≤ 100000, 0 ≤ Ai ≤ 100000, 1 ≤ U, V ≤ N\)

分析

樹形\(DP\)

我們設 \(f[i][j]\) 為以\(i\)作為根節點的子樹中公差為\(j\)的路徑的最長長度，接下來考慮轉移

轉移的過程無非是把子樹中的狀態遞迴至父親節點

即 \(f[now][a[now]-a[u]]=max(f[now][a[now]-a[u]],f[u][a[now]-a[u]]+1)\)

其中 \(now\) 為父親節點，\(u\)為兒子節點

統計答案時，我們只要在所有的\(f[now][val]+f[now][-val]+1\)

其實就是把兩條鏈拼在一起

加上一是為了防止特判一些奇奇怪怪的邊界問題，比如說只有一個點的情況

要注意 \(0\) 的時候要特判一下，因為此時\(val\)和\(-val\) 相等，直接更新會出錯

程式碼

#include<bits/stdc++.h>
using namespace std;
const int maxn=1e6+5;
struct asd{
	int from,to,next;
}b[maxn];
int head[maxn],tot=1,n,a[maxn],ans=0;
inline void ad(int aa,int bb){
	b[tot].from=aa;
	b[tot].to=bb;
	b[tot].next=head[aa];
	head[aa]=tot++;
}
inline int read(){
	register int x=0,f=1;
	char ch=getchar();
	while(ch<'0' || ch>'9'){
		if(ch=='-') f=-1;
		ch=getchar();
	}
	while(ch>='0' && ch<='9'){
		x=(x<<1)+(x<<3)+(ch^48);
		ch=getchar();
	}
	return x*f;
}
map<int,int>f[maxn];
void dfs(int now,int fa){
	int max0=0;
	for(int i=head[now];i!=-1;i=b[i].next){
		int u=b[i].to;
		if(u==fa) continue;
		dfs(u,now);
		if(a[u]==a[now]){
			if(f[now][0]<=f[u][0]+1){
				max0=f[now][0];
				f[now][0]=f[u][0]+1;
			}
			else if(max0<f[u][0]+1) max0=f[u][0]+1;
			ans=max(ans,f[now][0]+max0+1);
		} else {
			f[now][a[now]-a[u]]=max(f[now][a[now]-a[u]],f[u][a[now]-a[u]]+1);
			ans=max(ans,f[now][a[now]-a[u]]+f[now][a[u]-a[now]]+1);
		}
	}
}
int main(){
	freopen("C.in","r",stdin);
	freopen("C.out","w",stdout);
	memset(head,-1,sizeof(head));
	n=read();
	for(int i=1;i<=n;i++){
		a[i]=read();
	}
	for(register int i=1;i<n;i++){
		register int aa,bb;
		aa=read(),bb=read();
		ad(aa,bb);
		ad(bb,aa);
	}
	dfs(1,0);
	printf("%d\n",ans);
	return 0;
}