gym/102411 ICPC 2019-2020 North-Western Russia Regional Contest E. Equidistant

https://codeforces.com/gym/102411/problem/E

mx[u]表示 参赛团队所在城市到 u 的最远距离， mn[u]表示 参赛团队所在城市到 u 的最近距离， 以u为根通过树形dp可以得到mx[u]、mn[u]，再通过换根dp， 得到以其余点为根时的结果， 当mx[u] == mn[u]时， u即为答案。

换根时(由u换到v)， 提前保存 u 的子节点 中mx最大和次大， 当 mx[u] = mx[v] + 1时， 换根后的 mx[u] = 次大子节点 + 1。同理最小与次小也要保存。

#include
using namespace std;
typedef long long ll;
const int N = 2e5+10;
const int INF = 1e9+10;
int n, m;
vector g[N];
int c[N];
int mx[N], mn[N];
int ans = 0;
void dfs1(int u, int fa){
	for(auto v:g[u]){
		if(v == fa) continue;
		dfs1(v, u);
		mx[u] = max(mx[u], mx[v]+1);
		mn[u] = min(mn[u], mn[v]+1);
	}
}


void dfs2(int u, int fa){
	if(ans) return ;
	if(mx[u] == mn[u]){
		ans = u;
		return ;
	}
	int mx1[2] = {-INF, -INF};
	int mn1[2] = {INF, INF};
	if(c[u]){
		mx1[1] = mn1[1] = 0;
	}
	for(auto v:g[u]){
		if(mx[v]+1 > mx1[1]){
			mx1[0] = mx1[1];
			mx1[1] = mx[v]+1;
		}
		else if(mx[v]+1 > mx1[0]){
			mx1[0] = mx[v]+1;
		}
		if(mn[v]+1 < mn1[1]){
			mn1[0] = mn1[1];
			mn1[1] = mn[v]+1;
		}
		else if(mn[v]+1 < mn1[0]){
			mn1[0] = mn[v]+1;
		}
	}
	for(auto v:g[u]){
		if(v == fa) continue;
		int vmx = mx[v];
		int vmn = mn[v];
		if(mx[u] == mx[v]+1){
			mx[u] = mx1[0];
			mx[v] = max(mx[v], mx[u]+1);
		}
		else{
			mx[v] = max(mx[v], mx[u] + 1);
		}
		if(mn[u] == mn[v]+1){
			mn[u] = mn1[0];
			mn[v] = min(mn[v], mn[u]+1);
		}
		else{
			mn[v] = min(mn[v], mn[u] + 1);
		}
		dfs2(v, u);
		mx[u] = mx1[1];
		mn[u] = mn1[1];
		mx[v] = vmx;
		mn[v] = vmn;
	}
}

int main(){
	scanf("%d%d", &n, &m);
	int u, v;
	for(int i = 1; i < n; i++){
		scanf("%d%d", &u, &v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	for(int i = 1; i <= n; i++){
		mx[i] = -INF, mn[i] = INF;
	}
	int x;
	for(int i = 1; i <= m; i++){
		scanf("%d", &x);
		mx[x] = mn[x] = 0;
		c[x] = 1;
	}
	dfs1(1, 0);
	dfs2(1, 0);
	if(ans){
		printf("YES\n%d\n", ans);
	}
	else printf("NO\n");
	return 0;
}