树上的等差数列 HihoCoder - 1676 树形DP

题目链接：Hihocoder-1676

主要思路：

不难看出本题要用到树形DP，而边权差值有正有负，故dp数组第二维开map，dp[i][j]代表对于每个结点维护从这个结点走值为j的边最多能走dp[i][j]个点（不包括i）。dp[i][j]=max(dp[nxt][j]+1)，同时更新ans=max（ans,dp[now][j]+dp[now][-j]+1）。若val[i]==val[nxt]，那么若更新会将同一条链算两次，故在此时另开一数组特判，存最长的val值相等的两条链（DP即为从i出发最多向下可以连续走DP[i].d1个与其val值相同的点 （不包括i) ,次多为DP[i].d2）。

AC代码：

#include<cstdio>
#include<algorithm>
#include<map>
#define M 100005
using namespace std;
struct E {
	int to,nx;
} edge[M<<1];
int tot,head[M];
void Addedge(int a,int b) {
	edge[++tot].to=b;
	edge[tot].nx=head[a];
	head[a]=tot;
}
struct MX {
	int d1,d2;
	MX() {
		d1=d2=0;
	}
	void Add(int x) {
		if(x>d1) {
			d2=d1;
			d1=x;
		} else if(x>d2)d2=x;
	}
	int sum(){
		return d1+d2;
	}
} DP[M];//DP即为从i出发最多向下可以连续走DP[i]个与其val值相同的点 （不包括i) 
map<int,int>dp[M];//dp即为从i出发最多向下走dp[i][j]个与其父节点val值相差j的点 （不包括i)
int dis[M],ans=0;
int val[M];
void dfs(int now,int fa) {
	for(int i=head[now]; i; i=edge[i].nx) {
		int nxt=edge[i].to;
		if(nxt==fa)continue;
		dfs(nxt,now);
		if(val[now]==val[nxt]) {
			DP[now].Add(DP[nxt].d1+1);
			ans=max(ans,DP[now].sum()+1);
		} else {
			dp[now][val[now]-val[nxt]]=max(dp[nxt][val[now]-val[nxt]]+1,dp[now][val[now]-val[nxt]]);//更新dp值 
			ans=max(ans,dp[now][val[nxt]-val[now]]+dp[now][val[now]-val[nxt]]+1);
		}
		dp[nxt].clear();
	}
}
int main() {
	int n;
	scanf("%d",&n);
	for(int i=1; i<=n; i++)scanf("%d",&val[i]);
	for(int i=1; i<n; i++) {
		int a,b;
		scanf("%d%d",&a,&b);
		Addedge(a,b);
		Addedge(b,a);
	}
	dfs(1,-1);
	printf("%d\n",ans);
}