結果

問題 No.1002 Twotone
ユーザー tsutajtsutaj
提出日時 2020-02-21 01:40:10
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
TLE  
実行時間 -
コード長 5,513 bytes
コンパイル時間 858 ms
コンパイル使用メモリ 67,984 KB
実行使用メモリ 24,384 KB
最終ジャッジ日時 2024-10-10 01:41:45
合計ジャッジ時間 23,798 ms
ジャッジサーバーID
(参考情報)
judge4 / judge3
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
10,496 KB
testcase_01 AC 2 ms
5,248 KB
testcase_02 AC 3 ms
5,248 KB
testcase_03 AC 270 ms
12,160 KB
testcase_04 AC 277 ms
14,976 KB
testcase_05 AC 279 ms
15,104 KB
testcase_06 AC 2 ms
5,248 KB
testcase_07 AC 965 ms
14,864 KB
testcase_08 AC 2,170 ms
23,036 KB
testcase_09 AC 2,184 ms
21,836 KB
testcase_10 AC 2 ms
5,248 KB
testcase_11 AC 1,188 ms
20,524 KB
testcase_12 AC 1,598 ms
24,316 KB
testcase_13 AC 1,504 ms
24,236 KB
testcase_14 AC 3 ms
5,248 KB
testcase_15 AC 909 ms
14,580 KB
testcase_16 AC 2,176 ms
24,384 KB
testcase_17 AC 1,864 ms
24,360 KB
testcase_18 AC 3 ms
5,248 KB
testcase_19 TLE -
testcase_20 -- -
testcase_21 -- -
testcase_22 -- -
testcase_23 -- -
testcase_24 -- -
testcase_25 -- -
testcase_26 -- -
testcase_27 -- -
testcase_28 -- -
testcase_29 -- -
testcase_30 -- -
testcase_31 -- -
testcase_32 -- -
testcase_33 -- -
testcase_34 -- -
権限があれば一括ダウンロードができます

ソースコード

diff #

#include <cstdio>
#include <tuple>
#include <vector>
#include <utility>
#include <algorithm>
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);
    }

    long long int ans = 0, ans_dbl = 0;
    auto dfs = [&](auto&&f, int cur, int par) -> vector< pair<int, int> > {
        vector< vector< pair<int, int> > > children;
        vector< pair<int, int> > rec;
        for(auto e : G[cur]) {
            int to, color; tie(to, color) = e;
            if(to == par) continue;
            auto vec = f(f, to, cur);
            children.emplace_back();
            for(auto edge : vec) {
                if(edge.first == color or edge.second == color) {
                    children.back().emplace_back(edge);
                    rec.emplace_back(edge);
                    if(edge.first != -1 and edge.second != -1) {
                        ans++;
                    }
                }
                else if(edge.first == -1) {
                    edge.first = color;
                    if(edge.first > edge.second) swap(edge.first, edge.second);
                    children.back().emplace_back(edge);
                    rec.emplace_back(edge);
                    if(edge.first != -1 and edge.second != -1) {
                        ans++;
                    }
                }
            }

            // cur と to の間のパス
            {
                pair<int, int> short_edge(-1, color);
                children.back().emplace_back(short_edge);
                rec.emplace_back(short_edge);
            }
            
            sort(children.back().begin(), children.back().end());
        }
                
        sort(rec.begin(), rec.end());
        for(const auto &v : children) {
            for(const auto &e : v) {
                if(e.first != -1 and e.second != -1) {
                    // 完全に一致するもの (2 回数えられる)
                    {
                        auto ub = upper_bound(rec.begin(), rec.end(), e);
                        auto lb = lower_bound(rec.begin(), rec.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(rec.begin(), rec.end(), one);
                            auto lb = lower_bound(rec.begin(), rec.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(rec.begin(), rec.end(), one);
                            auto lb = lower_bound(rec.begin(), rec.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(rec.begin(), rec.end(), p1);
                        auto lb = lower_bound(rec.begin(), rec.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(rec.begin(), rec.end(), p);
                        auto lb = lower_bound(rec.begin(), rec.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;
                    }
                }
            }
        }
        return rec;
    };
    dfs(dfs, 0, -1);
    printf("%lld\n", ans + ans_dbl / 2);
    return 0;
}
0