結果

問題 No.1002 Twotone
コンテスト
ユーザー momohara
提出日時 2020-05-22 18:50:56
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 1,359 ms / 5,000 ms
コード長 5,095 bytes
コンパイル時間 2,109 ms
コンパイル使用メモリ 203,396 KB
実行使用メモリ 54,016 KB
最終ジャッジ日時 2024-10-05 06:51:37
合計ジャッジ時間 14,782 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;

#define rep(i,m,n) for(int (i)=(int)(m);i<(int)(n);++i)
#define rep2(i,m,n) for(int (i)=(int)(n)-1;i>=(int)(m);--i)
#define REP(i,n) rep(i,0,n)
#define REP2(i,n) rep2(i,0,n)
#define FOR(i,c) for(decltype((c).begin())i=(c).begin();i!=(c).end();++i)
#define all(hoge) (hoge).begin(),(hoge).end()
#define en '\n'
using ll = long long;
using ull = unsigned long long;
template <class T> using vec = vector<T>;
template <class T> using vvec = vector<vec<T>>;
typedef pair<ll, ll> P;
constexpr long long INF = 1LL << 60;
constexpr int INF_INT = 1 << 25;
constexpr long long MOD = (ll) 1e9 + 7;
//constexpr long long MOD = 998244353LL;
using ld=long double;
static const ld pi = 3.141592653589793L;
typedef vector<ll> Array;
typedef vector<Array> Matrix;


template<class T> inline bool chmin(T& a, T b) {
    if (a > b) {
        a = b;
        return true;
    }
    return false;
}
template<class T> inline bool chmax(T& a, T b) {
    if (a < b) {
        a = b;
        return true;
    }
    return false;
}

struct Edge {
	ll to, cap, rev;
	Edge(ll _to, ll _cap, ll _rev) {
	to = _to; cap = _cap; rev = _rev;
	}
};
using Edges = vector<Edge>;
using Graph = vector<Edges>;

void add_edge(Graph& G, ll from, ll to, ll cap, bool revFlag, ll revCap) {
	G[from].push_back(Edge(to, cap, (ll)G[to].size()));
	if (revFlag)G[to].push_back(Edge(from, revCap, (ll)G[from].size() - 1));
}


class CentroidDecomposition {
    Graph& g;
    ll n;
    vector<int> sz, dead;
public:
    CentroidDecomposition(Graph& _g) :g(_g),n(g.size()),sz(n, 1), dead(n, 0) {}

    ll szdfs(ll v, ll par) {
        sz[v] = 1;
        for (auto e : g[v])
            if (e.to != par && !dead[e.to]) sz[v] += szdfs(e.to, v);
        return sz[v];
    }

    void find(ll v, ll par, ll tmp, Array& cs) {
        ll ok = 1;
        for (auto e : g[v]) {
            if (e.to == par || dead[e.to]) continue;
            find(e.to, v, tmp, cs);
            ok &= (sz[e.to] <= tmp / 2);
        }
        ok &= (tmp - sz[v] <= tmp / 2);
        if (ok) cs.push_back(v);
    };

    vector<ll> build(ll r) {
        ll tmp = szdfs(r, -1);
        vector<ll> cs;
        find(r, -1, tmp, cs);
        return cs;
    }

    inline ll get_sz(ll v) { return sz[v]; }

    inline void disable(ll v) { dead[v] = 1; }

    inline void enable(ll v) { dead[v] = 0; }

    inline ll alive(ll v) { return !dead[v]; }
};

void solve(){
    ll n,k;
    cin>>n>>k;
    Graph g(n);
    REP(i,n-1){
        ll a,b,c;
        cin>>a>>b>>c;
        a--;b--;
        add_edge(g,a,b,c,true,c);
    }

    CentroidDecomposition cd(g);

    ll sum1=0, spre1=0;
    map<ll,ll> p1,pre1;
    map<pair<ll,ll>,ll> p2,pre2;
    map<ll,ll> p21,pre21;

    auto init =[&](){
        sum1=0;
        p1.clear();
        p2.clear();
        p21.clear();
    };
    auto initpre = [&](){
        spre1=0;
        pre1.clear();
        pre2.clear();
        pre21.clear();
    };
    
    auto dfs = [&](auto &&self, int v, int p, int c1,int c2)->void{
        if(c1&&c2){
            if(c1>c2) swap(c1,c2);
            p2[{c1,c2}]++;
            p21[c1]++;
            p21[c2]++;
        }else if(c1){
            p1[c1]++;
            sum1++;
        }

        for(auto e:g[v]){
            if(e.to==p || !cd.alive(e.to)) continue;
            if(c1&&c2){
                if(e.cap == c1 || e.cap==c2) self(self,e.to,v,c1,c2); 
            }else if(c1){
                if(e.cap == c1) self(self,e.to,v,c1,c2);
                else self(self,e.to,v,c1,e.cap);
            }
        }
    };

    queue<int> que;
    que.push(0);
    ll ans=0;
    while(que.size()){
        int v=que.front();que.pop();
        int r=cd.build(v)[0];//重心
        cd.disable(r);
        initpre();
        for(auto e:g[r]){
            if(cd.alive(e.to)){
                que.push(e.to);
                init();
                dfs(dfs, e.to, r, e.cap, 0);

                //count
                for(auto i:p1){
                    ans+=(spre1-pre1[i.first])*i.second;
                    ans+=(pre21[i.first])*i.second;
                }
                for(auto i:p2){
                    ans+=i.second;
                    ans+=pre2[i.first]*i.second;
                    ans+=pre1[i.first.first]*i.second;
                    ans+=pre1[i.first.second]*i.second;
                }
                //add
                spre1+=sum1;
                for(auto i:p1){
                    //cout<<i.first<<" "<<i.second<<en;
                    pre1[i.first]+=i.second;
                }
                for(auto i:p2){
                    //cout<<"{"<<i.first.first<<","<<i.first.second<<"} "<<i.second<<en;
                    pre2[i.first]+=i.second;
                }
                for(auto i:p21){
                    pre21[i.first]+=i.second;
                }
            }
        }

        //cout<<r<<" "<<ans<<en;
        //cout<<en;
        
        
    }
    cout<<ans<<en;

}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
	
	solve();
	//ll t;cin>>t;REP(i,t) solve();

    return 0;
}
0