結果

問題 No.1002 Twotone
ユーザー tsutajtsutaj
提出日時 2020-02-21 01:37:58
言語 C++14
(gcc 12.3.0 + boost 1.83.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
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
6,816 KB
testcase_01 AC 2 ms
6,816 KB
testcase_02 AC 1 ms
6,820 KB
testcase_03 AC 240 ms
13,424 KB
testcase_04 AC 335 ms
16,676 KB
testcase_05 AC 344 ms
16,700 KB
testcase_06 AC 2 ms
6,816 KB
testcase_07 AC 360 ms
15,024 KB
testcase_08 AC 674 ms
22,684 KB
testcase_09 AC 683 ms
22,548 KB
testcase_10 AC 2 ms
6,816 KB
testcase_11 AC 676 ms
18,344 KB
testcase_12 AC 947 ms
23,032 KB
testcase_13 AC 936 ms
22,772 KB
testcase_14 AC 2 ms
6,820 KB
testcase_15 AC 522 ms
14,136 KB
testcase_16 AC 1,082 ms
22,584 KB
testcase_17 AC 1,094 ms
22,904 KB
testcase_18 AC 2 ms
6,816 KB
testcase_19 AC 587 ms
17,668 KB
testcase_20 AC 711 ms
23,048 KB
testcase_21 AC 698 ms
22,620 KB
testcase_22 AC 3 ms
6,816 KB
testcase_23 AC 805 ms
17,292 KB
testcase_24 AC 1,518 ms
27,444 KB
testcase_25 AC 1,478 ms
26,372 KB
testcase_26 AC 2 ms
6,816 KB
testcase_27 AC 104 ms
23,612 KB
testcase_28 AC 164 ms
32,356 KB
testcase_29 AC 157 ms
31,432 KB
testcase_30 AC 2 ms
6,816 KB
testcase_31 AC 162 ms
31,748 KB
testcase_32 AC 174 ms
31,480 KB
testcase_33 AC 159 ms
31,776 KB
testcase_34 AC 1,122 ms
26,668 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#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;
}
0