Code mẫu

#include <bits/stdc++.h>
using namespace std;
#define tag "LITPATH"
#define maxn 100007
#define forinc(i,a,b) for(int i=a;i<=b;i++)
#define checkfile(FiLeNaMe) { if(fopen(FiLeNaMe".inp","r")) freopen(FiLeNaMe".inp","r",stdin),freopen(FiLeNaMe".out","w",stdout); }
/// my data:
using unlong = unsigned long long;
#define long long long
struct neigh{
    int ver,wei;
};
///>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
int n;
int64_t ans;
vector<neigh> adj[maxn];
int f[maxn][4];
///0 --> 00
///1 --> 11
///2 --> 10
///3 --> 01

void dfs(int u=1,int p=0){

    for(auto nei: adj[u]){
        int v=nei.ver;
        int w=nei.wei;

        if(v==p) continue;

        dfs(v,u);

        ++f[v][w];

        if(w){
            ans+=f[v][1]<<1;
            ans+=f[v][0];
            ans+=f[v][2];

            ans+=(int64_t)f[u][0]*f[v][1];
            ans+=(int64_t)f[u][1]*f[v][0];
            ans+=(int64_t)f[u][1]*f[v][1]<<1;

            ans+=(int64_t)f[u][1]*f[v][2];
            ans+=(int64_t)f[u][2]*f[v][1];

            f[u][1]+=f[v][1];
            f[u][2]+=f[v][2]+f[v][0];
        }else{

            ans+=(int64_t)(f[u][0]+1)*f[v][1];
            ans+=(int64_t)(f[u][0]+1)*f[v][3];
            ans+=(int64_t)(f[u][0]+1)*f[v][0]<<1;

            ans+=(int64_t)f[u][1]*f[v][0];
            ans+=(int64_t)f[u][3]*f[v][0];

            f[u][0]+=f[v][0];
            f[u][3]+=f[v][3]+f[v][1];
        }
    }
}
void enter(){
    cin>>n;
    int u,v,w;
    forinc(i,2,n) cin>>u>>v>>w,adj[u].push_back({v,w}),adj[v].push_back({u,w});
}

void solve(){
    dfs();
    cout<<ans<<"\n";
}
int main()
{
    checkfile(tag);
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    enter();
    solve();

    return 0;
}