金题大战Vol.0 C、树上的等差数列

金题大战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
",ans);
	return 0;
}