結果

問題 No.1002 Twotone
コンテスト
ユーザー ミドリムシ
提出日時 2020-02-28 23:24:09
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
WA  
実行時間 -
コード長 6,431 bytes
コンパイル時間 2,030 ms
コンパイル使用メモリ 197,740 KB
実行使用メモリ 26,456 KB
最終ジャッジ日時 2024-10-13 19:14:12
合計ジャッジ時間 19,826 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 13 WA * 4 TLE * 1 -- * 15
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;
using lint = long long;
const lint mod = 1e9 + 7;
#define all(x) (x).begin(), (x).end()
#define bitcount(n) __builtin_popcountl((lint)(n))
#define fcout cout << fixed << setprecision(15)
#define highest(x) (63 - __builtin_clzl(x))
template<class T> inline void YES(T condition){ if(condition) cout << "YES" << endl; else cout << "NO" << endl; }
template<class T> inline void Yes(T condition){ if(condition) cout << "Yes" << endl; else cout << "No" << endl; }
template<class T = string, class U = char>int character_count(T text, U character){ int ans = 0; for(U i: text){ ans += (i == character); } return ans; }
lint power(lint base, lint exponent, lint module){ if(exponent % 2){ return power(base, exponent - 1, module) * base % module; }else if(exponent){ lint root_ans = power(base, exponent / 2, module); return root_ans * root_ans % module; }else{ return 1; }}
struct position{ int y, x; }; position mv[4] = {{0, -1}, {1, 0}, {0, 1}, {-1, 0}}; // double euclidean(position first, position second){ return sqrt((second.x - first.x) * (second.x - first.x) + (second.y - first.y) * (second.y - first.y)); }
template<class T, class U> string to_string(pair<T, U> x){ return to_string(x.first) + "," + to_string(x.second); } string to_string(string x){ return x; }
template<class itr> void array_output(itr start, itr goal){ string ans; for(auto i = start; i != goal; i++) ans += to_string(*i) + " "; if(!ans.empty()) ans.pop_back(); cout << ans << endl; }
template<class itr> void cins(itr first, itr last){ for(auto i = first; i != last; i++){ cin >> (*i); } }
template<class T> T gcd(T a, T b){ if(a && b){ return gcd(min(a, b), max(a, b) % min(a, b)); }else{ return a; }} template<class T> T lcm(T a, T b){ return a / gcd(a, b) * b; }
struct combination{ vector<lint> fact, inv; combination(int sz) : fact(sz + 1), inv(sz + 1){ fact[0] = 1; for(int i = 1; i <= sz; i++){ fact[i] = fact[i - 1] * i % mod; } inv[sz] = power(fact[sz], mod - 2, mod); for(int i = sz - 1; i >= 0; i--){ inv[i] = inv[i + 1] * (i + 1) % mod; } } lint C(int p, int q) const{ if(q < 0 || p < q) return 0; return (fact[p] * inv[q] % mod * inv[p - q] % mod); } };
template<class itr> bool next_sequence(itr first, itr last, int max_bound){ itr now = last; while(now != first){ now--; (*now)++; if((*now) == max_bound){ (*now) = 0; }else{ return true; } } return false; }

struct road{
    int to, color;
};

struct colorpair{
    int c1, c2;
};

colorpair merge(colorpair p1, colorpair p2){
    if(p1.c1 == -2 || p2.c1 == -2){
        return {-2, -2};
    }
    if(p1.c1 == p2.c1 && p1.c2 == p2.c2){
        return p1;
    }
    if(p1.c2 != -1 && p2.c2 != -1){
        return {-2, -2};
    }
    if(p1.c2 == -1 && p2.c2 == -1){
        return {p1.c1, p2.c1};
    }
    if(p1.c1 == -1){
        swap(p1, p2);
    }
    if(p1.c1 == p2.c1 || p1.c2 == p2.c1){
        return p1;
    }
    return {-2, -2};
}

int N, K;
vector<vector<road>> roads;
vector<bool> is_in_forest;


vector<int> subtree_size;

int get_subtree_size(int now, int back){
    //cout << now << " ";
    int ret = 1;
    for(road i: roads[now]){
        if(i.to != back && is_in_forest[i.to]){
            ret += get_subtree_size(i.to, now);
        }
    }
    return subtree_size[now] = ret;
}

int full_size;

int find_centroid(int now, int back){
    int max_subtree = -1, max_subtree_size = 0;
    for(road i: roads[now]){
        if(i.to != back && is_in_forest[i.to]){
            if(max_subtree_size < subtree_size[i.to]){
                max_subtree_size = subtree_size[i.to];
                max_subtree = i.to;
            }
        }
    }
    if(max_subtree_size * 2 <= full_size){
        return now;
    }else{
        return find_centroid(max_subtree, now);
    }
}

vector<vector<colorpair>> colorpairs;

void colorful_dfs(int now, int back, int number, colorpair now_color){
    if(now_color.c1 == -2){
        return;
    }
    colorpairs[number].push_back(now_color);
    for(road i: roads[now]){
        if(i.to != back && is_in_forest[i.to]){
            colorful_dfs(i.to, now, number, merge(now_color, {i.color, -1}));
        }
    }
}

lint centroid_decomposition(int root){
    //cout << root << endl;
    get_subtree_size(root, -1);
    //cout << endl;
    full_size = subtree_size[root];
    int centroid = find_centroid(root, -1);
    colorpairs.clear();
    vector<road> subroot;
    for(road i: roads[root]){
        if(is_in_forest[i.to]){
            subroot.push_back(i);
        }
    }
    for(int i = 0; i < subroot.size(); i++){
        colorpairs.push_back({});
        colorful_dfs(subroot[i].to, root, i, {subroot[i].color, -1});
        /*for(auto j: colorpairs[i]){
         cout << j.c1 << "," << j.c2 << " ";
         }
         cout << endl;*/
    }
    //cout << endl;
    
    map<pair<int, int>, lint> color_all, color_subtree[subroot.size()];
    lint alone_all = 0, alone_subtree[subroot.size()];
    fill(alone_subtree, alone_subtree + subroot.size(), 0);
    for(int i = 0; i < subroot.size(); i++){
        for(auto j: colorpairs[i]){
            color_all[{j.c1, j.c2}]++;
            color_subtree[i][{j.c1, j.c2}]++;
            if(j.c2 == -1){
                alone_all++;
                alone_subtree[i]++;
            }
        }
    }
    lint ans = 0, half = 0;
    for(int i = 0; i < subroot.size(); i++){
        for(auto j: colorpairs[i]){
            if(j.c2 != -1){
                half += color_all[{j.c1, j.c2}] - color_subtree[i][{j.c1, j.c2}];
                ans += color_all[{j.c1, -1}] - color_subtree[i][{j.c1, -1}];
                ans += color_all[{j.c2, -1}] - color_subtree[i][{j.c2, -1}];
                ans += 1;
            }else{
                half += (alone_all - alone_subtree[i]) - (color_all[{j.c1, -1}] - color_subtree[i][{j.c1, -1}]);
            }
        }
    }
    assert(half % 2 == 0);
    ans += half / 2;
    
    is_in_forest[root] = false;
    for(road i: roads[root]){
        if(is_in_forest[i.to]){
            ans += centroid_decomposition(i.to);
        }
    }
    return ans;
}


int main(){
    cin >> N >> K;
    roads.resize(N);
    for(int i = 0; i < N - 1; i++){
        int u, v, c;
        cin >> u >> v >> c;
        u--, v--, c--;
        roads[u].push_back({v, c});
        roads[v].push_back({u, c});
    }
    is_in_forest.resize(N, true);
    subtree_size.resize(N);
    cout << centroid_decomposition(0) << endl;
}
0