結果

問題 No.1002 Twotone
ユーザー tsutaj
提出日時 2020-02-21 01:40:10
言語 C++14
(gcc 13.3.0 + boost 1.87.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
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 17 TLE * 1 -- * 15
権限があれば一括ダウンロードができます

ソースコード

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