1156D - 0-1-Tree - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
#pragma GCC optimize ("Ofast")
#pragma GCC target ("avx2")
#define el << '\n'
#define nl cout << '\n'
#define cina(a,n) for (int i=0; i<n; i++) cin >> a[i]
#define ll long long
#define spc << ' '
#define forn(i,n) for (int i=0; i<n; i++)
#define forl(i,s,e) for (int i=s; i<e; i++)
#define MOD 1000000007
#define MOD2 998244353
#define pb push_back
#define ansyes cout << "YES\n"
#define ansno cout << "NO\n"
#define fi first
#define se second
#define pii pair<int,int>
#define pll pair<long long, long long>

int n;
vector<pair<int, bool>> acc[200000];
pll dp[200000];
ll tot=0;

void dfs1(int cn, int pn) {
	dp[cn].fi = 1;
	// if (acc[cn].size()>1) dp[cn].se = 1;
	for (pair<int, bool> i : acc[cn]) {
		if (i.fi == pn) continue;
		// if (cn == 6) cout << i.fi el;
		dfs1(i.fi, cn);
		if (i.se) {
			dp[cn].se += dp[i.fi].se;
			dp[cn].se += dp[i.fi].fi;
		}
		else {
			dp[cn].fi += dp[i.fi].fi;
		}
	}
	return;
}

void reroot(int cr, int nr, bool ct) {
	// if ()
	if (!ct) {
		dp[cr].fi -= dp[nr].fi;
		dp[nr].fi += dp[cr].fi;
	}
	else {
		dp[cr].se -= dp[nr].fi;
		dp[cr].se -= dp[nr].se;
		dp[nr].se += dp[cr].fi;
		dp[nr].se += dp[cr].se;
	}
	return;
}

void dfs2(int cn, int pn, bool ct) {
	tot += dp[cn].fi + dp[cn].se - 1;
	// if (cn == 3) {
		// forn(i,n) cout << dp[i].fi spc << dp[i].se el;
	// }
	for (pair<int, bool> i : acc[cn]) {
		if (i.fi == pn) continue;
		reroot(cn, i.fi, i.se);
		dfs2(i.fi, cn, i.se);
	}
	if (pn != -1) reroot(cn, pn, ct);
}

int main() {
	cin >> n;
	forn(i,n-1) {
		pair<int, bool> t1, t2; cin >> t1.fi >> t2.fi;
		t1.fi--;
		t2.fi--;
		bool ct; cin >> ct;
		if (ct==1) {
			t1.se = true;
			t2.se = true;
		}
		else {
			t1.se = false;
			t2.se = false;
		}
		acc[t1.fi].pb(t2);
		acc[t2.fi].pb(t1);
	}
	dfs1(0, -1);
	// forn(i,n) cout << dp[i].fi spc << dp[i].se el;
	dfs2(0, -1, 0);
	cout << tot el;
}