結果

問題 No.1002 Twotone
ユーザー outlineoutline
提出日時 2020-03-03 22:48:42
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
WA  
実行時間 -
コード長 4,156 bytes
コンパイル時間 1,581 ms
コンパイル使用メモリ 133,644 KB
実行使用メモリ 139,392 KB
最終ジャッジ日時 2024-10-13 23:19:57
合計ジャッジ時間 20,245 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge4
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 4 ms
7,936 KB
testcase_01 AC 4 ms
8,192 KB
testcase_02 AC 4 ms
7,936 KB
testcase_03 WA -
testcase_04 WA -
testcase_05 WA -
testcase_06 AC 4 ms
7,936 KB
testcase_07 WA -
testcase_08 WA -
testcase_09 WA -
testcase_10 AC 4 ms
8,044 KB
testcase_11 WA -
testcase_12 WA -
testcase_13 WA -
testcase_14 WA -
testcase_15 WA -
testcase_16 WA -
testcase_17 WA -
testcase_18 AC 5 ms
7,936 KB
testcase_19 WA -
testcase_20 WA -
testcase_21 WA -
testcase_22 WA -
testcase_23 WA -
testcase_24 WA -
testcase_25 WA -
testcase_26 AC 4 ms
8,192 KB
testcase_27 AC 119 ms
44,272 KB
testcase_28 AC 198 ms
54,276 KB
testcase_29 AC 186 ms
54,156 KB
testcase_30 AC 4 ms
8,064 KB
testcase_31 AC 178 ms
53,668 KB
testcase_32 AC 191 ms
53,712 KB
testcase_33 AC 183 ms
53,840 KB
testcase_34 WA -
権限があれば一括ダウンロードができます

ソースコード

diff # 

#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;
            ans += it.second * it.second;
        }
    }

    for(auto it : cnt){
        for(auto it2 : it){
            if(it2.first.size() == 1)   ;
            else{
                twosum[it2.first] -= it2.second * it2.second;
                ans -= it2.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