【树的点分治】阴阳

阴阳

【题目大意】有时间再写

【题解】。。

#include<cstdio>
#include<cstring>
#include<algorithm>
#define fo(i,a,b) for (int i = a;i <= b;i ++)

using namespace std;
typedef long long LL;

const int P = 100000;
const int maxn = 100005;
int N,tot,root,size,cnt;
int a[maxn],b[maxn*2],c[maxn*2],d[maxn*2];
int f[maxn],s[maxn];
bool flag[maxn];
LL ans;

struct T
{
	int t;
	LL n;
	inline LL Q(int TIME)
	{
		if (TIME != t) n = 0;
		t = TIME;
		return n;
	}
};

void Insert(int x,int y,int z)
{
	tot ++;
	b[tot] = y;
	c[tot] = a[x];
	a[x] = tot;
	d[tot] = z;
}

void Findroot(int x,int last)
{
	s[x] = 1;
	f[x] = 0;
	for (int i = a[x];i;i = c[i])
		if (b[i] != last && !flag[b[i]])
		{
			Findroot(b[i],x);
			s[x] += s[b[i]];
			f[x] = max(f[x],s[b[i]]);
		}
	f[x] = max(f[x],size-s[x]);
	if (f[x] < f[root]) root = x;
}

T tank[2][maxn*2];
int cur[maxn*2];

void dfs(int x,int last,int D,int TIME)
{
	if (cur[D+P] > 1)
	{
		ans += tank[0][-D+P].Q(TIME);
		if (cur[D+P] > 2 && !D) ans ++;
	}
	else ans += tank[1][-D+P].Q(TIME);
	for (int i = a[x];i;i = c[i])
		if (b[i] != last && !flag[b[i]])
		{
			if (!d[i])
			{
				cur[D+1+P] ++;
				dfs(b[i],x,D+1,TIME);
				cur[D+1+P] --;
			}
			else
			{
				cur[D-1+P] ++;
				dfs(b[i],x,D-1,TIME);
				cur[D-1+P] --;
			}
		}
}

inline void update(T &x,int TIME)
{
	if (x.t == TIME) x.n ++;
	else
	{
		x.t = TIME;
		x.n = 1;
	}
}

void collect(int x,int last,int D,int TIME)
{
	update(tank[0][D+P],TIME);
	if (cur[D+P] > 1) update(tank[1][D+P],TIME);
	for (int i = a[x];i;i = c[i])
		if (b[i] != last && !flag[b[i]])
		{
			if (!d[i])
			{
				cur[D+1+P] ++;
				collect(b[i],x,D+1,TIME);
				cur[D+1+P] --;
			}
			else
			{
				cur[D-1+P] ++;
				collect(b[i],x,D-1,TIME);
				cur[D-1+P] --;
			}
		}
}

void Work(int x,int TIME)
{
	cur[P] = 1;
	for (int i = a[x];i;i = c[i])
		if (!flag[b[i]])
		{
			cur[(!d[i]?1:-1)+P] ++;
			dfs(b[i],x,!d[i]?1:-1,TIME);
			collect(b[i],x,!d[i]?1:-1,TIME);
			cur[(!d[i]?1:-1)+P] --;
		}
	flag[x] = 1;
	for (int i = a[x];i;i = c[i])
		if (!flag[b[i]])
		{
			f[0] = size = s[b[i]];
			Findroot(b[i],root = 0);
			Work(root,++cnt);
		}
}

int main()
{
	freopen("t.in","r",stdin);
	freopen("t.out","w",stdout);
	scanf("%d",&N);
	fo(i,1,N-1)
	{
		int x,y,z;
		scanf("%d%d%d",&x,&y,&z);
		Insert(x,y,z);
		Insert(y,x,z);
	}
	f[0] = size = N;
	Findroot(1,root = 0);
	Work(root,cnt = 1);
	printf("%I64d\n",ans);
	return 0;
}