树上的等差数列

题目描述

给定一棵包含 $N$ 个节点的无根树，节点编号 $1\to 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

样例输出

数据范围与提示

对于 $50\%$ 的数据，$1 \leqslant N \leqslant 1000$

对于 $100\%$ 的数据，$1 \leqslant N \leqslant 100000, 0 \leqslant A_i \leqslant 100000, 1 \leqslant U, V \leqslant N$

分析

树形 $dp$ 好题。

因为要求的是最长的等差序列，根节点不同，答案也可能不同，所以 $dp$ 的状态转移就定义为 $f[i][j]$ 表示 $i$ 节点为根，公差为 $j$ 时的最长的等差数列，不包括自己。那么我们就可以愉快的 $dfs$ 来进行转移了。

我们记录一下他自己和他的父亲，避免出现死循环，每一次先 $dfs$ 到儿子，递归上来，然后就处理出来了公差为 $\Delta$ 的以儿子为根的所有长度，这时候我们只需要判断一下此时的 $\Delta$ 值是否为 $0$。如果是，那么 $ans$ 的转移应该是：

\[ans = max(ans,f[x][0] + f[son[x]][0] + 2)
\]

因为此时 $f[x][0]$ 存储的是其他儿子上最长链，所以需要加上当前儿子的最长链，因为我们的数组不保存自己，所以要加 $2$ 。

其他情况就是直接更新 $ans$ ，他的答案应该是 $f[x][d] + f[x][-d] + 1$ ，因为他的父亲那里也可能会有链，公差为 $-d$ 就是那个链，由于负数下标的问题，我们利用 $map$ 来存储，然后轻松解决此题。

代码

#include<cstdio>

#include<cstring>

#include<iostream>

#include<algorithm>

#include<map>

#define re register

using namespace std;

const int maxn = 1e5+10;

map <int,int> mp[maxn];

struct Node{

	int v,next;

}e[maxn<<1];

int w[maxn];

int ans = 0;

int head[maxn],tot;

void Add(int x,int y){//建边

	e[++tot].v = y;

	e[tot].next = head[x];

	head[x] = tot;

}

inline int read(){//快读

	int s = 0,f = 1;

	char ch = getchar();

	while(!isdigit(ch)){if(ch=='-')f=-1;ch=getchar();}

	while(isdigit(ch)){s=s*10+ch-'0';ch=getchar();}

	return s * f;

}

inline void DP(int x,int fa){

	for(int i=head[x];i;i=e[i].next){

		int v = e[i].v;

		if(v == fa)continue;//避免死循环

		int d = w[v] - w[x];//计算公差

		DP(v,x);

		if(!d){//公差为0的情况

			ans = max(ans,mp[x][0] + mp[v][0] + 2);

			mp[x][0] = max(mp[x][0],mp[v][0] + 1);

		}

		else{//公差不为0

			mp[x][d] = max(mp[x][d],mp[v][d] + 1);

			ans = max(ans,mp[x][d] + mp[x][-d] + 1);

		}

	}

}

int main(){

	freopen("C.in","r",stdin);

	freopen("C.out","w",stdout);

	int n =read();

	for(re int i = 1;i<=n;++i){w[i]=read();}

	for(re int i = 1;i< n;++i){

		int x = read(),y = read();

		Add(x,y);

		Add(y,x);

	}

	DP(1,0);

	printf("%d\n",ans);

}