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