結果

問題 No.1002 Twotone
ユーザー tsutajtsutaj
提出日時 2020-02-21 01:37:58
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 1,499 ms / 5,000 ms
コード長 6,575 bytes
コンパイル時間 1,083 ms
コンパイル使用メモリ 76,740 KB
実行使用メモリ 31,936 KB
最終ジャッジ日時 2023-07-31 07:45:49
合計ジャッジ時間 18,671 ms
ジャッジサーバーID
(参考情報)
judge11 / judge14
このコードへのチャレンジ(β)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
4,380 KB
testcase_01 AC 1 ms
4,380 KB
testcase_02 AC 1 ms
4,380 KB
testcase_03 AC 240 ms
13,092 KB
testcase_04 AC 331 ms
16,252 KB
testcase_05 AC 326 ms
16,196 KB
testcase_06 AC 1 ms
4,380 KB
testcase_07 AC 347 ms
14,848 KB
testcase_08 AC 643 ms
22,404 KB
testcase_09 AC 655 ms
22,332 KB
testcase_10 AC 1 ms
4,376 KB
testcase_11 AC 645 ms
18,120 KB
testcase_12 AC 897 ms
22,800 KB
testcase_13 AC 927 ms
22,508 KB
testcase_14 AC 1 ms
4,376 KB
testcase_15 AC 536 ms
13,744 KB
testcase_16 AC 1,088 ms
22,572 KB
testcase_17 AC 1,086 ms
22,460 KB
testcase_18 AC 1 ms
4,376 KB
testcase_19 AC 606 ms
17,392 KB
testcase_20 AC 739 ms
22,744 KB
testcase_21 AC 738 ms
22,388 KB
testcase_22 AC 2 ms
4,376 KB
testcase_23 AC 805 ms
16,904 KB
testcase_24 AC 1,499 ms
27,036 KB
testcase_25 AC 1,475 ms
25,848 KB
testcase_26 AC 1 ms
4,376 KB
testcase_27 AC 102 ms
23,300 KB
testcase_28 AC 159 ms
31,936 KB
testcase_29 AC 159 ms
31,320 KB
testcase_30 AC 1 ms
4,376 KB
testcase_31 AC 155 ms
31,776 KB
testcase_32 AC 169 ms
31,148 KB
testcase_33 AC 159 ms
31,104 KB
testcase_34 AC 1,099 ms
26,448 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