洛谷 [P3629] 巡逻

树的直径

树的直径有两种求法
1.两遍 dfs 法, 便于输出具体方案,但是无法处理负权边
2.DP 法,代码量少,可以处理负权边

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <cmath>
#include <cstring>
#include <cstdlib>
using namespace std;
const int MAXN = 200005;
int n, k, head[MAXN], nume, fa[MAXN], ans, ma, t, d[MAXN], s;
bool f[MAXN];
struct edge {
	int to, nxt, dis;
}e[MAXN<<1];
void adde(int from, int to) {
	e[++nume].to = to;
	e[nume].dis = 1;
	e[nume].nxt = head[from];
	head[from] = nume;
}
void dfs1(int u, int dep) {
	if(f[u]) return ;
	f[u] = 1;
	if(dep > ma && u != 1) {
		ma = dep;
		k = u;
	}
	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		dfs1(v, dep + e[i].dis);
	}
}
void del(int u) {
	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if(v == fa[u]) {
			e[i].dis = e[((i - 1) ^ 1) + 1].dis = -1;
			del(v);
		}
	}
}
void dfs2(int u, int dep) {
	f[u] = 1;
	if(dep > ma && u != t) {
		ma = dep; k = u;
	}

	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if(!f[v]) {
			fa[v] = u;
			dfs2(v, dep + e[i].dis);
		}
	}
}
void dp(int u) {
	f[u] = 1;
	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if(!f[v]) {
			dp(v);
			ma = max(ma, d[u] + d[v] + e[i].dis);
			d[u] = max(d[u], d[v] + e[i].dis);
		}
	}
}
int main() {
	cin >> n >> s;
	for(int i = 1; i < n; i++) {
		int u, v;
		scanf("%d %d", &u, &v);
		adde(u, v);
		adde(v, u);
	}
	dfs1(1, 0);
	ma = 0;t = k; k = 0;
	memset(f, 0, sizeof(f));
	dfs2(t, 0);
	ans += 2 * (n - 1) ;
	ans -= ma - 1;
	if(s == 2) {
		del(k);
		memset(f, 0, sizeof(f));
		ma = 0;
		dp(1);
		ans -= ma - 1;
	}
	cout << ans << endl;
	return 0;
}