結果

問題 No.1002 Twotone
コンテスト
ユーザー sugarrr
提出日時 2020-03-01 14:56:17
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 1,390 ms / 5,000 ms
コード長 3,237 bytes
コンパイル時間 2,383 ms
コンパイル使用メモリ 194,256 KB
実行使用メモリ 70,624 KB
最終ジャッジ日時 2024-10-13 20:41:33
合計ジャッジ時間 18,705 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

//#include <bits/stdc++.h>
#include "bits/stdc++.h"
using namespace std;
typedef long long ll;
//#include "boost/multiprecision/cpp_int.hpp"
//typedef boost::multiprecision::cpp_int ll;
typedef long double dd;
#define i_7 (ll)(1E9+7)
//#define i_7 998244353
#define i_5 i_7-2
ll mod(ll a){
    ll c=a%i_7;
    if(c>=0)return c;
    return c+i_7;
}
typedef pair<ll,ll> l_l;
typedef pair<dd,dd> d_d;
ll inf=(ll)1E16;
#define rep(i,l,r) for(ll i=l;i<=r;i++)
#define pb push_back
ll max(ll a,ll b){if(a<b)return b;else return a;}
ll min(ll a,ll b){if(a>b)return b;else return a;}
void Max(ll &pos,ll val){pos=max(pos,val);}//Max(dp[n],dp[n-1]);
void Min(ll &pos,ll val){pos=min(pos,val);}
void Add(ll &pos,ll val){pos=mod(pos+val);}
dd EPS=1E-9;
#define fastio ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define fi first
#define se second
////////////////////////////

ll n,k;
#define N 200005
vector<l_l>v[N];
ll sz[N];
bool used[N];
ll szdfs(ll x,ll p){
    sz[x]=1;
    for(auto z:v[x]){
        ll y=z.fi;
        if(y==p||used[y])continue;
        sz[x]+=szdfs(y,x);
    }
    return sz[x];
}
ll centroid(ll x,ll p,ll total){
    for(auto z:v[x]){
        ll y=z.fi;
        if(y==p||used[y])continue;
        if(2*sz[y]>total)return centroid(y,x,total);
    }
    return x;
}
/////

bool over[N];
set<ll>col[N];
ll ans=0;
ll calc(vector<ll>t){
    ll res=0;
    map<ll,ll>s1,s2;
    ll c1=0,c2=0;
    map<l_l,ll>same;
    bool zero=false;
    for(auto u:t){
        if(col[u].size()==1){
            for(auto x:col[u]){
                s1[x]++;
            }
            c1++;
        }else if(col[u].size()==2){
            ll z[2];ll i=0;
            for(auto x:col[u]){
                s2[x]++;
                z[i]=x;i++;
            }
            same[l_l(z[0],z[1])]++;
            c2++;
        }else if(col[u].size()==0){
            zero=true;
        }
    }
    if(zero)res=c2*2;
    for(auto u:s1){
        res+=u.se*(c1-u.se);
        res+=2*u.se*s2[u.fi];
    }
    for(auto u:same){
        res+=u.se*(u.se-1);
    }
    return res;
}
void dfs(vector<ll>&t,ll x,ll p){
    t.pb(x);
    for(auto z:v[x]){
        ll y=z.fi;ll c=z.se;
        if(y==p||used[y])continue;
        if(over[x]){
            over[y]=true;continue;
        }
        //col[y].clear();
        col[y]=col[x];col[y].insert(c);
        if(col[y].size()>2){
            over[y]=true;
            continue;
        }
        over[y]=false;
        dfs(t,y,x);
    }
}
void solve(ll x,ll p){
    szdfs(x,p);
    x=centroid(x, p, sz[x]);
    used[x]=true;
    vector<ll>T;T.pb(x);
    col[x].clear();
    over[x]=false;
    for(auto z:v[x]){
        ll y=z.fi;ll c=z.se;
        vector<ll>t;
        if(y==p||used[y])continue;
        over[y]=false;
        col[y].clear();
        col[y].insert(c);
        dfs(t,y,x);
        ans-=calc(t);
        T.insert(T.end(),t.begin(),t.end());
    }
    ans+=calc(T);
    for(auto z:v[x]){
        ll y=z.fi;
        if(y==p||used[y])continue;
        solve(y,x);
    }
}
int main(){fastio
    cin>>n>>k;
    rep(i,1,n-1){
        ll a,b,c;cin>>a>>b>>c;
        v[a].pb(l_l(b,c));
        v[b].pb(l_l(a,c));
    }
    solve(1,-1);
    cout<<ans/2<<endl;
    return 0;
}
/*
4 0
1 2 1
2 3 2
3 4 3
*/
0