結果

問題 No.1002 Twotone
コンテスト
ユーザー outline
提出日時 2020-03-03 22:09:23
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 3,527 ms / 5,000 ms
コード長 4,595 bytes
コンパイル時間 1,644 ms
コンパイル使用メモリ 136,092 KB
実行使用メモリ 139,520 KB
最終ジャッジ日時 2024-10-13 23:09:12
合計ジャッジ時間 22,710 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 33
権限があれば一括ダウンロードができます

ソースコード

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