結果

問題 No.1002 Twotone
ユーザー momoharamomohara
提出日時 2020-05-22 18:09:54
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
WA  
実行時間 -
コード長 5,164 bytes
コンパイル時間 3,255 ms
コンパイル使用メモリ 197,776 KB
実行使用メモリ 56,448 KB
最終ジャッジ日時 2024-04-15 10:28:56
合計ジャッジ時間 32,787 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
5,248 KB
testcase_01 AC 2 ms
5,376 KB
testcase_02 AC 2 ms
5,376 KB
testcase_03 WA -
testcase_04 WA -
testcase_05 WA -
testcase_06 AC 3 ms
5,376 KB
testcase_07 AC 465 ms
15,232 KB
testcase_08 WA -
testcase_09 WA -
testcase_10 AC 2 ms
5,376 KB
testcase_11 WA -
testcase_12 WA -
testcase_13 WA -
testcase_14 AC 2 ms
5,376 KB
testcase_15 AC 692 ms
15,488 KB
testcase_16 AC 1,485 ms
24,576 KB
testcase_17 AC 1,400 ms
24,696 KB
testcase_18 AC 2 ms
5,376 KB
testcase_19 AC 1,581 ms
25,216 KB
testcase_20 AC 1,927 ms
32,512 KB
testcase_21 AC 1,933 ms
32,000 KB
testcase_22 AC 2 ms
5,376 KB
testcase_23 AC 1,303 ms
25,472 KB
testcase_24 AC 2,377 ms
38,528 KB
testcase_25 AC 2,354 ms
38,528 KB
testcase_26 AC 2 ms
5,376 KB
testcase_27 AC 154 ms
16,640 KB
testcase_28 AC 270 ms
21,556 KB
testcase_29 AC 225 ms
21,496 KB
testcase_30 AC 2 ms
5,376 KB
testcase_31 AC 207 ms
21,024 KB
testcase_32 AC 297 ms
21,432 KB
testcase_33 AC 229 ms
21,304 KB
testcase_34 AC 2,495 ms
56,448 KB
権限があれば一括ダウンロードができます

ソースコード

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);

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

    auto init =[&](){
        mp.clear();
        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 x){
        mp[x]++;
        if(mp.size()==1){
            auto it=mp.begin();
            p1[it->first]++;
            sum1++;
        }else if(mp.size()==2){
            auto it=mp.begin();
            ll a=it->first;
            it++;
            ll b=it->first;
            p2[{a,b}]++;
            p21[a]++;
            p21[b]++;
        }else{
            return;
        }
        if(cd.alive(v)){
            for(auto e:g[v]){
                if(e.to==p) continue;
                self(self, e.to, v, e.cap);
            }
        }

        mp[x]--;
        if(mp[x]==0) mp.erase(x);
    };

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

                //count
                for(auto i:p1){
                    ans+=(spre1-pre1[i.first])*i.second;
                    ans+=(pre21[i.first])*i.second;
                }
                for(auto i:p2){
                    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;
        cd.disable(r);
        
    }
    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