#432147 (C++14) No.1002 Twotone

• 問題: No.1002 Twotone
• ユーザー: tsutaj
• 提出日時: 2020-02-21 01:37:58
• 言語: C++14 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 1,518 ms / 5,000 ms
• コード長: 6,575 bytes
• コンパイル時間: 1,477 ms
• コンパイル使用メモリ: 77,296 KB
• 実行使用メモリ: 32,356 KB
• 最終ジャッジ日時: 2024-10-10 01:24:00
• 合計ジャッジ時間: 18,569 ms
• ジャッジサーバーID （参考情報）: judge4 / judge5

ファイルパターン	結果
sample	AC * 2
other	AC * 33

ソースコード

#include <cstdio>
#include <vector>
#include <algorithm>
#include <tuple>
#include <utility>
using namespace std;
using Graph = vector< vector< pair<int, int> > >;

int main() {
    int N, K; scanf("%d%d", &N, &K);
    Graph G(N);
    for(int i=0; i+1<N; i++) {
        int u, v, c; scanf("%d%d%d", &u, &v, &c);
        u--; v--; c--;
        G[u].emplace_back(v, c);
        G[v].emplace_back(u, c);
    }

    vector<int> num(N), dead(N);

    auto get_size = [&](auto&& f, int cur, int par) -> void {
        num[cur] = 1;
        for(auto e : G[cur]) {
            int to = e.first;
            if(to == par or dead[to]) continue;
            f(f, to, cur);
            num[cur] += num[to];
        }
    };

    auto find_centroid = [&](int root) {
        get_size(get_size, root, -1);
        int V = num[root];
        auto dfs = [&](auto&& f, int cur, int par) -> int {
            for(auto e : G[cur]) {
                int to = e.first;
                if(to == par or dead[to]) continue;
                if(num[to] > V / 2) return f(f, to, cur);
            }
            return cur;
        };
        return dfs(dfs, root, -1);
    };

    auto get_edges = [&](int cur)
        -> vector< vector< pair<int, int> > > {
        vector< vector< pair<int, int> > > res;

        auto dfs = [&](auto &&f, int cur, int par, pair<int, int> edge) -> void {
            for(auto e : G[cur]) {
                pair<int, int> tmp = edge;
                int to, color; tie(to, color) = e;
                if(to == par or dead[to]) continue;
                if(tmp.first == color or tmp.second == color) {
                    res.back().emplace_back(tmp);
                    f(f, to, cur, tmp);
                }
                else if(tmp.first == -1) {
                    tmp.first = color;
                    if(tmp.first > tmp.second) swap(tmp.first, tmp.second);
                    res.back().emplace_back(tmp);
                    f(f, to, cur, tmp);
                }
            }
        };

        for(auto e : G[cur]) {
            int to, color; tie(to, color) = e;
            if(dead[to]) continue;
            res.emplace_back();
            res.back().emplace_back(-1, color);
            dfs(dfs, to, cur, make_pair(-1, color));
            sort(res.back().begin(), res.back().end());
        }
        return res;
    };


    long long int ans = 0, ans_dbl = 0;
    auto rec = [&](auto &&f, int root) -> void {
        int c = find_centroid(root);
        dead[c] = true;

        for(auto e : G[c]) {
            int to = e.first;
            if(dead[to]) continue;
            f(f, to);
        }

        // c を通るものを数える
        auto vec = get_edges(c);
        vector< pair<int, int> > flat;
        for(auto v : vec) for(auto e : v) flat.emplace_back(e);
        sort(flat.begin(), flat.end());

        for(const auto &v : vec) {
            for(const auto &e : v) {
                if(e.first != -1 and e.second != -1) {
                    // それ単体でも答えになりうる
                    ans++;
                    
                    // 完全に一致するもの (2 回数えられる)
                    {
                        auto ub = upper_bound(flat.begin(), flat.end(), e);
                        auto lb = lower_bound(flat.begin(), flat.end(), e);
                        ans_dbl += ub - lb;
                    }
                    {
                        auto ub = upper_bound(v.begin(), v.end(), e);
                        auto lb = lower_bound(v.begin(), v.end(), e);
                        ans_dbl -= ub - lb;
                    }
                    
                    // first だけ
                    {
                        pair<int, int> one(-1, e.first);
                        {
                            auto ub = upper_bound(flat.begin(), flat.end(), one);
                            auto lb = lower_bound(flat.begin(), flat.end(), one);
                            ans += ub - lb;
                        }
                        {
                            auto ub = upper_bound(v.begin(), v.end(), one);
                            auto lb = lower_bound(v.begin(), v.end(), one);
                            ans -= ub - lb;
                        }
                    }
                    
                    // second だけ
                    {
                        pair<int, int> one(-1, e.second);
                        {
                            auto ub = upper_bound(flat.begin(), flat.end(), one);
                            auto lb = lower_bound(flat.begin(), flat.end(), one);
                            ans += ub - lb;
                        }
                        {
                            auto ub = upper_bound(v.begin(), v.end(), one);
                            auto lb = lower_bound(v.begin(), v.end(), one);
                            ans -= ub - lb;
                        }                    
                    }
                }
                else {
                    // 一色だけなら、他の一色だけのものと合わせられるかも
                    int c = e.second;
                    {
                        pair<int, int> p1(-1, 1 << 28);
                        pair<int, int> p2(-1, -1);
                        auto ub = upper_bound(flat.begin(), flat.end(), p1);
                        auto lb = lower_bound(flat.begin(), flat.end(), p2);
                        ans_dbl += ub - lb;
                    }
                    {
                        pair<int, int> p1(-1, 1 << 28);
                        pair<int, int> p2(-1, -1);
                        auto ub = upper_bound(v.begin(), v.end(), p1);
                        auto lb = lower_bound(v.begin(), v.end(), p2);
                        ans_dbl -= ub - lb;
                    }
                    {
                        pair<int, int> p(-1, c);
                        auto ub = upper_bound(flat.begin(), flat.end(), p);
                        auto lb = lower_bound(flat.begin(), flat.end(), p);
                        ans_dbl -= ub - lb;
                    }
                    {
                        pair<int, int> p(-1, c);
                        auto ub = upper_bound(v.begin(), v.end(), p);
                        auto lb = lower_bound(v.begin(), v.end(), p);
                        ans_dbl += ub - lb;
                    }
                }                
            }
        }

        dead[c] = false;
    };

    rec(rec, 0);
    printf("%lld\n", ans + ans_dbl / 2);
    return 0;
}