HDU 5016 Mart Master II

Trader Dogy lives in city S, which consists of n districts. There are n - 1 bidirectional roads in city S, each connects a pair of districts. Indeed, city S is connected, i.e. people can travel between every pair of districts by roads. 

In some districts there are marts founded by Dogy’s competitors. when people go to marts, they’ll choose the nearest one. In cases there are more than one nearest marts, they’ll choose the one with minimal city number. 

Dogy’s money could support him to build only one new marts, he wants to attract as many people as possible, that is, to build his marts in some way that maximize the number of people who will choose his mart as favorite. Could you help him?

There are multiple test cases. Please process till EOF. 

In each test case: 

First line: an integer n indicating the number of districts. 

Next n - 1 lines: each contains three numbers b _i, e _i and w _i, (1 ≤ b _i,e _i ≤ n,1 ≤ w _i ≤ 10000), indicates that there’s one road connecting city b _i and e _i, and its length is w _i. 

Last line : n(1 ≤ n ≤ 10 ⁵) numbers, each number is either 0 or 1, i-th number is 1 indicates that the i-th district has mart in the beginning and vice versa.

For each test case, output one number, denotes the number of people you can attract, taking district as a unit.

5
1 2 1
2 3 1
3 4 1
4 5 1
1 0 0 0 1
5
1 2 1
2 3 1
3 4 1
4 5 1
1 0 0 0 0
1
1
1
0 

2
4
0
1 

给一棵树，每个为0的点会被最近的标号小的为1的点覆盖，现在要把一个值为0的点变成1，问变1后这个点最多能覆盖多少为0的点。先用最短路算出每个点被谁覆盖，然后用树分治统计每个点变1能覆盖多少点。
#include<queue>
#include<cstdio>
#include<vector>
#include<iostream>
#include<algorithm>
using namespace std;
const int INF = 0x7FFFFFFF;
const int maxn = 2e5 + 10;
int n, x, y, z, ans[maxn], flag[maxn];

struct Tree
{
	int ft[maxn], nt[maxn], u[maxn], v[maxn], sz;
	int mx[maxn], ct[maxn], vis[maxn];
	int dis[maxn], id[maxn];
	struct point
	{
		int x, y, z;
		point(int x = 0, int y = 0, int z = 0) :x(x), y(y), z(z) {}
		bool operator<(const point&a)const { return y == a.y ? z > a.z : y > a.y; }
	};
	point d[maxn], D[maxn];
	int t;
	void clear(int n)
	{
		mx[sz = 0] = INF;
		for (int i = 1; i <= n; i++) dis[i] = ft[i] = -1, ans[i] = vis[i] = 0;
	}
	void AddEdge(int x, int y, int z)
	{
		u[sz] = y;	v[sz] = z;	nt[sz] = ft[x];	ft[x] = sz++;
		u[sz] = x;	v[sz] = z;	nt[sz] = ft[y]; ft[y] = sz++;
	}
	void dijkstra()
	{
		priority_queue<point> p;
		for (int i = 1; i <= n; i++)
		{
			scanf("%d", &flag[i]);
			if (flag[i]) p.push(point(i, dis[i] = 0, id[i] = i));
		}
		while (!p.empty())
		{
			point q = p.top();	p.pop();
			if (vis[q.x]) continue; vis[q.x] = 1;
			for (int i = ft[q.x]; i != -1; i = nt[i])
			{
				if (!vis[q.x]) continue;
				if (dis[u[i]] == -1 || dis[u[i]] > q.y + v[i])
				{
					dis[u[i]] = q.y + v[i];
					id[u[i]] = q.z;
					p.push(point(u[i], dis[u[i]], q.z));
				}
				else if (dis[u[i]] == q.y + v[i] && id[u[i]] > q.z)
				{
					id[u[i]] = q.z;
					p.push(point(u[i], dis[u[i]], q.z));
				}
			}
		}
		for (int i = 1; i <= n; i++)
		{
			vis[i] = 0;
			if (dis[i] == -1) dis[i] = INF;
		}
	}
	int dfs(int x, int fa, int sum)
	{
		int y = mx[x] = (ct[x] = 1) - 1;
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (vis[u[i]] || u[i] == fa) continue;
			int z = dfs(u[i], x, sum);
			ct[x] += ct[u[i]];
			mx[x] = max(mx[x], ct[u[i]]);
			y = mx[y] < mx[z] ? y : z;
		}
		mx[x] = max(mx[x], sum - ct[x]);
		return mx[x] < mx[y] ? x : y;
	}
	void get(int x, int fa, int len)
	{
		d[t++] = point(x, len);
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (vis[u[i]] || u[i] == fa) continue;
			get(u[i], x, len + v[i]);
		}
	}
	void find(int x, int k, int len)
	{
		t = 0;  get(x, -1, len);
		for (int i = 0; i < t; i++) D[i] = point(0, dis[d[i].x] - d[i].y, id[d[i].x]);
		sort(D, D + t);
		for (int i = 0,j; i < t; i++)
		{
			j = lower_bound(D, D + t, point(0, d[i].y, d[i].x)) - D;
			ans[d[i].x] += k*j;
		}
	}
	void work(int x, int sum)
	{
		int y = dfs(x, -1, sum);
		find(y, 1, 0);	vis[y] = 1;
		for (int i = ft[y]; i != -1; i = nt[i])
		{
			if (vis[u[i]]) continue;
			find(u[i], -1, v[i]);
			work(u[i], ct[u[i]] > ct[y] ? sum - ct[y] : ct[u[i]]);
		}
	}
}solve;

int main()
{
	while (scanf("%d", &n) != EOF)
	{
		solve.clear(n);
		for (int i = 1; i < n; i++)
		{
			scanf("%d%d%d", &x, &y, &z);
			solve.AddEdge(x, y, z);
		}
		solve.dijkstra();
		solve.work(1, n);
		int res = 0;
		for (int i = 1; i <= n; i++) if (!flag[i]) res = max(res, ans[i]);
		printf("%d\n", res);
	}
	return 0;
}