比赛 2025.6.2 评测结果 AAAAAAAAAA
题目名称 0-1-Tree 最终得分 100
用户昵称 彭欣越 运行时间 1.348 s
代码语言 C++ 内存使用 24.29 MiB
提交时间 2025-06-02 12:01:45
显示代码纯文本
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=200010;
ll n,mk[N][5],cnt;
ll head[N],tot;
long long ans;
vector<int>v0[N],v1[N];
struct edge {
	int v,w,nxt;
}e[N*2];
void add (int u,int v,int w) {
	e[++tot].v=v;
	e[tot].w=w;
	e[tot].nxt=head[u];
	head[u]=tot; 
}
void dfs (int u,int fa,int k) {
	if (k) v1[cnt].push_back(u);
	else v0[cnt].push_back(u);
	mk[u][k]=cnt;
	for (int i=head[u];i;i=e[i].nxt) {
		int v=e[i].v,w=e[i].w;
	    if (w!=k||v==fa) continue;
	    dfs(v,u,k);
	}
}
int main () {
	freopen("0-1-Tree.in","r",stdin);
	freopen("0-1-Tree.out","w",stdout);
	cin >> n;
	for (int i=1;i<n;i++) {
		int a,b,c;
		cin >> a >> b >> c;
		add(a,b,c);
		add(b,a,c);
	}
	for (int i=1;i<=n;i++) {
		if (!mk[i][0]) {
		    cnt++;
		    dfs(i,0,0);
		}
		if (!mk[i][1]) {
		    cnt++;
		    dfs(i,0,1);
		}
	}
	for (int i=1;i<=cnt;i++) {
		ll l1=v0[i].size(),l2=v1[i].size();
		ans+=l1*(l1-1)+l2*(l2-1);
	}
	for (int i=1;i<=n;i++) {
		ll l1=v0[mk[i][0]].size(),l2=v1[mk[i][1]].size();
		ans+=(l1-1)*(l2-1);
	}
	cout << ans <<endl;
	return 0;
}