View Code of Problem 1072

#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long LL;
const double pi = acos(-1.0);
const int low(int x) { return x&-x; }
const int INF = 0x7FFFFFFF;
const int mod = 1e9 + 7;
const int maxn = 1e5 + 10;
int n, m, x, y, z, K;
char ch[5];

struct Tree
{
	const static int maxn = 1e5 + 10;
	int ft[maxn], nt[maxn], u[maxn], v[maxn], sz, tot;
	int dis[maxn], vis[maxn], cnt[maxn], mx[maxn];
	void clear(int n)
	{
		mx[sz = 0] = INF;
		for (int i = 1; i <= n; i++)
		{
			ft[i] = -1;	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++;
	}
	int dfs(int x, int fa, int sum)
	{
		int z = mx[x] = 0;
		cnt[x] = 1;
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (u[i] == fa || vis[u[i]]) continue;
			int y = dfs(u[i], x, sum);
			cnt[x] += cnt[u[i]];
			mx[x] = max(mx[x], cnt[u[i]]);
			if (mx[z] > mx[y]) z = y;
 		}
		mx[x] = max(mx[x], sum - cnt[x]);
		return mx[x] < mx[z] ? x : z;
	}
	void dfsdis(int x, int fa, int len)
	{
		dis[tot++] = len;
		for (int i = ft[x]; i != -1; i = nt[i])
		{
			if (u[i] == fa || vis[u[i]]) continue;
			dfsdis(u[i], x, len + v[i]);
		}
	}
	LL find(int x, int d)
	{
		tot = 0;
		dfsdis(x, -1, d);
		sort(dis, dis + tot);
		LL ans = 0;
		for (int i = 0, j = tot - 1; i < j; i++)
		{
			while (j > i&&dis[i] + dis[j] > K) j--;
			ans += j - i;
		}
		return ans;
	}
	LL work(int x, int sum)
	{
		int y = dfs(x, -1, sum);
		LL ans = find(y, 0);	vis[y] = 1;
		for (int i = ft[y]; i != -1; i = nt[i])
		{
			if (vis[u[i]]) continue;
			ans -= find(u[i], v[i]);
			ans += work(u[i], cnt[u[i]] > cnt[y] ? sum - cnt[y] : cnt[u[i]]);
		}
		return ans;
	}
}solve;

int main()
{
	while (scanf("%d%d", &n, &m) != EOF)
	{
		solve.clear(n);
		while (m--)
		{
			scanf("%d%d%d%s", &x, &y, &z, ch);
			solve.AddEdge(x, y, z);
		}
		scanf("%d", &K);
		printf("%d\n", solve.work(1, n));
	}
	return 0;
}

Double click to view unformatted code.


Back to problem 1072