結果

問題 No.1002 Twotone
ユーザー tsutajtsutaj
提出日時 2020-02-21 01:37:58
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 1,253 ms / 5,000 ms
コード長 6,575 bytes
コンパイル時間 1,307 ms
コンパイル使用メモリ 76,856 KB
実行使用メモリ 33,148 KB
最終ジャッジ日時 2024-04-18 08:12:22
合計ジャッジ時間 16,290 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
6,816 KB
testcase_01 AC 1 ms
6,940 KB
testcase_02 AC 1 ms
6,940 KB
testcase_03 AC 204 ms
13,832 KB
testcase_04 AC 294 ms
17,020 KB
testcase_05 AC 282 ms
17,012 KB
testcase_06 AC 2 ms
6,940 KB
testcase_07 AC 326 ms
15,668 KB
testcase_08 AC 637 ms
23,068 KB
testcase_09 AC 612 ms
23,060 KB
testcase_10 AC 2 ms
6,944 KB
testcase_11 AC 600 ms
18,840 KB
testcase_12 AC 811 ms
23,544 KB
testcase_13 AC 843 ms
23,232 KB
testcase_14 AC 2 ms
6,940 KB
testcase_15 AC 482 ms
14,648 KB
testcase_16 AC 961 ms
23,096 KB
testcase_17 AC 958 ms
23,412 KB
testcase_18 AC 2 ms
6,940 KB
testcase_19 AC 494 ms
18,144 KB
testcase_20 AC 580 ms
23,572 KB
testcase_21 AC 613 ms
23,180 KB
testcase_22 AC 2 ms
6,944 KB
testcase_23 AC 673 ms
17,932 KB
testcase_24 AC 1,253 ms
27,696 KB
testcase_25 AC 1,250 ms
26,884 KB
testcase_26 AC 2 ms
6,940 KB
testcase_27 AC 95 ms
23,116 KB
testcase_28 AC 152 ms
31,944 KB
testcase_29 AC 146 ms
32,172 KB
testcase_30 AC 2 ms
6,944 KB
testcase_31 AC 138 ms
31,888 KB
testcase_32 AC 154 ms
33,148 KB
testcase_33 AC 136 ms
31,908 KB
testcase_34 AC 979 ms
27,432 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