HDU 5886 Tower Defence（2016青岛网络赛 I题，树的直径 + DP）

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b)	for (int i(a); i <= (b); ++i)
#define dec(i, a, b)	for (int i(a); i >= (b); --i)
#define fi		first
#define se		second

typedef long long LL;

const int N = 100010;

vector <pair<int, LL > > v[N];
int T, n, L, R, x, cnt;
int a[N], f[N], father[N];
LL b[N], c[N], s[N], cl[N], cr[N], c1, c2, num, ans = 0;

void dfs1(int x, int fa, LL now){
	if (now > c1){
		c1 = now;
		L = x;
	}

	for (auto node : v[x]){
		int u = node.fi;
		LL w = node.se;
		if (u == fa) continue;
		dfs1(u, x, now + w);
	}
}

void dfs2(int x, int fa, LL now){
	father[x] = fa;
	s[x] = now;
	if (now > c2){
		c2 = now;
		R = x;
	}

	for (auto node : v[x]){
		int u = node.fi;
		LL w = node.se;
		if (u == fa) continue;
		dfs2(u, x, now + w);
	}
}

void dfs3(int w, int x, LL now){
	f[x] = 1;
	c[w] = max(c[w], now);
	for (auto node : v[x]){
		int u = node.fi;
		if (f[u]) continue;
		dfs3(w, u, now + node.se);
	}
}

int main(){

	for (scanf("%d", &T); T--; ){
		scanf("%d", &n);
		rep(i, 0, n + 1) v[i].clear();
		rep(i, 2, n){
			int x, y;
			LL z;
			scanf("%d%d%lld", &x, &y, &z);
			v[x].push_back({y, z});
			v[y].push_back({x, z});
		}

		L = -1; c1 = -1;
		dfs1(1, 0, 0);
		R = -1, c2 = -1;
		memset(father, 0, sizeof father);
		dfs2(L, 0, 0);

		x = R;
		cnt = 0;
		while (true){
			a[++cnt] = R;
			R = father[R];
			if (R == 0) break;
		}

		rep(i, 1, cnt) b[i] = s[a[i]];
		reverse(a + 1, a + cnt + 1);
		reverse(b + 1, b + cnt + 1);
		
		num = b[cnt];
		ans = num * (n - cnt);

		memset(f, 0, sizeof f);
		rep(i, 1, cnt) f[a[i]] = 1;

		rep(i, 1, cnt){
			c[i] = 0;
			dfs3(i, a[i], 0);
		}


		cl[1] = b[1]; rep(i, 2, cnt) cl[i] = max(cl[i - 1], b[i] + c[i]);
		cr[cnt] = 0;  dec(i, cnt - 1, 1) cr[i] = max(cr[i + 1], b[cnt] - b[i] + c[i]);
		rep(i, 1, cnt - 1) ans = ans + max(cl[i], cr[i + 1]);
		printf("%lld\n", ans);
	}

	return 0;
}