結果

問題 No.1002 Twotone
ユーザー outlineoutline
提出日時 2020-03-03 22:06:59
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 3,713 ms / 5,000 ms
コード長 4,662 bytes
コンパイル時間 1,683 ms
コンパイル使用メモリ 136,068 KB
実行使用メモリ 139,392 KB
最終ジャッジ日時 2024-10-13 23:08:37
合計ジャッジ時間 22,768 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge4
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 5 ms
8,064 KB
testcase_01 AC 4 ms
8,192 KB
testcase_02 AC 5 ms
8,064 KB
testcase_03 AC 498 ms
14,848 KB
testcase_04 AC 662 ms
16,896 KB
testcase_05 AC 652 ms
16,896 KB
testcase_06 AC 5 ms
7,936 KB
testcase_07 AC 290 ms
13,184 KB
testcase_08 AC 542 ms
16,512 KB
testcase_09 AC 546 ms
16,512 KB
testcase_10 AC 5 ms
8,064 KB
testcase_11 AC 682 ms
16,768 KB
testcase_12 AC 929 ms
18,944 KB
testcase_13 AC 937 ms
19,072 KB
testcase_14 AC 5 ms
8,192 KB
testcase_15 AC 498 ms
13,312 KB
testcase_16 AC 987 ms
17,280 KB
testcase_17 AC 964 ms
17,280 KB
testcase_18 AC 6 ms
8,192 KB
testcase_19 AC 671 ms
23,808 KB
testcase_20 AC 847 ms
33,920 KB
testcase_21 AC 861 ms
33,152 KB
testcase_22 AC 6 ms
8,064 KB
testcase_23 AC 826 ms
36,992 KB
testcase_24 AC 1,563 ms
54,528 KB
testcase_25 AC 1,559 ms
54,400 KB
testcase_26 AC 5 ms
8,064 KB
testcase_27 AC 129 ms
44,364 KB
testcase_28 AC 215 ms
54,112 KB
testcase_29 AC 199 ms
54,156 KB
testcase_30 AC 5 ms
8,064 KB
testcase_31 AC 193 ms
53,800 KB
testcase_32 AC 203 ms
53,844 KB
testcase_33 AC 193 ms
53,836 KB
testcase_34 AC 3,713 ms
139,392 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

// 参考文献 : http://kmjp.hatenablog.jp/entry/2020/02/29/0930

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <queue>
#include <string>
#include <map>
#include <set>
#include <tuple>
#include <deque>
#include <numeric>
#include <bitset>
#include <iomanip>
#include <cassert>
#include <chrono>
#include <random>
#include <limits>
#include <iterator>
#include <functional>
#include <sstream>
#include <complex>
using namespace std;

typedef long long ll;
typedef pair<int, int> P;
typedef pair<int, double> Pid;
typedef pair<double, int> Pdi;
typedef pair<ll, int> Pl;
typedef pair<int, pair<int, int>> PP;
const double PI = 3.1415926535897932;   // acos(-1)
const double EPS = 1e-15;
const int INF = 1001001001;
const int mod = 1e+9 + 7;

#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define chadd(x, y) x = (x + y) % mod

int n, k;
vector<P> graph[200005];    // (to, color)
bool used[200005];          // 分割点として使われたかどうか
int subtree_size[200005];   // 部分木のサイズ

// 部分木サイズの計算
void calc_subtreesize(int u, int parent = -1){
    subtree_size[u] = 1;
    for(int i = 0; i < graph[u].size(); ++i){
        int v = graph[u][i].first;
        if(v == parent || used[v])  continue;
        calc_subtreesize(v, u);
        subtree_size[u] += subtree_size[v];
    }
}

// 重心を探索 (v が重心候補)
// 返り値は (部分木の最大サイズ, 重心) のペア
P search_centroid(int from, int sz, int parent = -1){
    P res = P(INF, -1);
    int sum = 1, maxsz = 0;
    for(int i = 0; i < graph[from].size(); ++i){
        int to = graph[from][i].first, color = graph[from][i].second;
        if(to == parent || used[to])    continue;
        
        chmin(res, search_centroid(to, sz, from));

        chmax(maxsz, subtree_size[to]);
        sum += subtree_size[to];
    }
    chmax(maxsz, sz - sum);
    chmin(res, P(maxsz, from));
    return res;
}

ll ans = 0;

void Count(int from, map<set<int>, ll> &cnt, set<int> table, int parent = -1){
    ++cnt[table];
    for(int i = 0; i < graph[from].size(); ++i){
        int to = graph[from][i].first, color = graph[from][i].second;
        if(to == parent || used[to])    continue;
        set<int> table2 = table;
        table2.insert(color);
        if(table2.size() >= 3)  continue;
        Count(to, cnt, table2, from);
    }
}

void solve(int from){
    calc_subtreesize(from);
    if(subtree_size[from] == 1) return;
    P info = search_centroid(from, subtree_size[from]);
    int center = info.second;

    vector<map<set<int>, ll>> cnt;
    map<set<int>, ll> total;
    for(int i = 0; i < graph[center].size(); ++i){
        int to = graph[center][i].first, color = graph[center][i].second;
        if(used[to])    continue;
        cnt.push_back({});
        Count(to, cnt.back(), {color}, center);
        for(auto it : cnt.back()){
            total[it.first] += it.second;
        }
    }

    ll onenum = 0, oneans = 0;
    map<set<int>, ll> twosum;
    for(auto it : total){
        if(it.first.size() == 1){
            onenum += it.second;
            oneans -= it.second * it.second;
        }
        else{   // it.first.size() == 2
            twosum[it.first] += it.second * it.second;
            ans += it.second;
            for(auto it2 : it.first){
                set<int> st = {it2};
                auto ite = total.lower_bound(st);
                if((*ite).first == st){
                    ans += (*ite).second * it.second;
                }
            }
        }
    }

    for(auto it : cnt){
        for(auto it2 : it){
            if(it2.first.size() == 1)   ;
            else{
                twosum[it2.first] -= it2.second * it2.second;
                for(auto it3 : it2.first){
                    set<int> st = {it3};
                    auto ite = it.lower_bound(st);
                    if((*ite).first == st){
                        ans -= (*ite).second * it2.second;
                    }
                }
            }
        }
    }
    
    oneans += onenum * onenum;
    ans += oneans / 2;
    for(auto it : twosum){
        ans += it.second / 2;
    }
    used[center] = true;

    for(int i = 0; i < graph[center].size(); ++i){
        int to = graph[center][i].first;
        if(used[to])    continue;
        solve(to);
    }
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n >> k;
    for(int i = 0; i < n - 1; ++i){
        int u, v, c;
        cin >> u >> v >> c;
        graph[--u].push_back(P(--v, c));
        graph[v].push_back(P(u, c));
    }
    solve(0);
    cout << ans << endl;
}
0