結果

問題 No.1002 Twotone
ユーザー nono00nono00
提出日時 2024-02-27 03:56:37
言語 C++23
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 848 ms / 5,000 ms
コード長 7,324 bytes
コンパイル時間 1,927 ms
コンパイル使用メモリ 146,892 KB
実行使用メモリ 33,780 KB
最終ジャッジ日時 2024-09-29 11:53:54
合計ジャッジ時間 13,970 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
5,248 KB
testcase_01 AC 2 ms
5,248 KB
testcase_02 AC 1 ms
5,248 KB
testcase_03 AC 213 ms
11,092 KB
testcase_04 AC 307 ms
13,416 KB
testcase_05 AC 289 ms
13,544 KB
testcase_06 AC 2 ms
5,248 KB
testcase_07 AC 145 ms
9,848 KB
testcase_08 AC 274 ms
14,184 KB
testcase_09 AC 278 ms
14,312 KB
testcase_10 AC 2 ms
5,248 KB
testcase_11 AC 284 ms
11,100 KB
testcase_12 AC 413 ms
13,800 KB
testcase_13 AC 409 ms
13,800 KB
testcase_14 AC 2 ms
6,820 KB
testcase_15 AC 198 ms
9,344 KB
testcase_16 AC 433 ms
13,800 KB
testcase_17 AC 409 ms
13,932 KB
testcase_18 AC 2 ms
6,820 KB
testcase_19 AC 457 ms
16,192 KB
testcase_20 AC 534 ms
21,604 KB
testcase_21 AC 586 ms
21,224 KB
testcase_22 AC 2 ms
6,820 KB
testcase_23 AC 362 ms
16,636 KB
testcase_24 AC 786 ms
24,904 KB
testcase_25 AC 848 ms
24,836 KB
testcase_26 AC 2 ms
6,816 KB
testcase_27 AC 57 ms
10,512 KB
testcase_28 AC 90 ms
14,444 KB
testcase_29 AC 86 ms
14,436 KB
testcase_30 AC 2 ms
6,820 KB
testcase_31 AC 79 ms
14,280 KB
testcase_32 AC 84 ms
14,436 KB
testcase_33 AC 82 ms
14,440 KB
testcase_34 AC 837 ms
33,780 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <iostream>
#include <map>
#include <vector>

#include <iterator>
#include <ranges>
#include <vector>

namespace nono {

template <class T>
struct EdgeBase {
    int from;
    int to;
    T weight;
    EdgeBase() {}
    EdgeBase(int from, int to, T weight = 1): from(from), to(to), weight(weight) {}
};

using Edge = EdgeBase<int>;
template <class T>
using WeightedEdge = EdgeBase<T>;

template <class T>
class Graph {
    struct Edge_ {
        int to;
        T weight;
        int id;
    };

    using iterator = std::vector<Edge_>::iterator;
    using const_iterator = std::vector<Edge_>::const_iterator;
    using subrange = std::ranges::subrange<iterator, iterator>;
    using const_subrange = std::ranges::subrange<const_iterator, const_iterator>;

  public:
    template <class U>
    friend Graph<U> to_undirected_graph(int n, const std::vector<EdgeBase<U>>& edges);
    template <class U>
    friend Graph<U> to_directed_graph(int n, const std::vector<EdgeBase<U>>& edges);

    subrange operator[](int i) {
        return std::ranges::subrange(edges_.begin() + indptr_[i], edges_.begin() + indptr_[i + 1]);
    }
    const_subrange operator[](int i) const {
        return std::ranges::subrange(edges_.begin() + indptr_[i], edges_.begin() + indptr_[i + 1]);
    }

    int size() const {
        return n_;
    }
    bool is_directed() const {
        return directed_;
    }
    bool is_undirected() const {
        return !is_directed();
    }

  private:
    Graph(int n, const std::vector<EdgeBase<T>>& edges, bool directed)
        : n_(n),
          indptr_(n_ + 1),
          edges_(directed ? edges.size() : 2 * edges.size()),
          directed_(directed) {
        for (const auto& e: edges) {
            indptr_[e.from + 1]++;
            if (!directed_) indptr_[e.to + 1]++;
        }
        for (int i = 0; i < n_; i++) {
            indptr_[i + 1] += indptr_[i];
        }
        auto index = indptr_;
        for (int i = 0; i < std::ssize(edges); i++) {
            const auto& e = edges[i];
            edges_[index[e.from]++] = Edge_(e.to, e.weight, i);
            if (!directed_) edges_[index[e.to]++] = Edge_(e.from, e.weight, i);
        }
    }

    int n_;
    std::vector<int> indptr_;
    std::vector<Edge_> edges_;
    bool directed_;
};

template <class T>
Graph<T> to_undirected_graph(int n, const std::vector<EdgeBase<T>>& edges) {
    return Graph<T>(n, edges, false);
}

template <class T>
Graph<T> to_directed_graph(int n, const std::vector<EdgeBase<T>>& edges) {
    return Graph<T>(n, edges, true);
}

}  //  namespace nono

#include <vector>

#include <algorithm>
#include <vector>

namespace nono {

template <class T>
bool is_tree(const Graph<T>& graph) {
    if (graph.is_directed()) return false;
    constexpr int NONE = -1;
    int n = graph.size();
    std::vector<int> used(n);
    auto dfs = [&](auto self, int u, int eid) -> bool {
        used[u] = 1;
        for (const auto& e: graph[u]) {
            if (e.id == eid) continue;
            if (used[e.to]) return false;
            if (!self(self, e.to, e.id)) return false;
        }
        return true;
    };
    if (!dfs(dfs, 0, NONE)) return false;
    return std::ranges::all_of(used, [](int f) { return f == 1; });
}

}  //  namespace nono

namespace nono {

template <class T>
std::vector<int> centroids(const Graph<T>& graph) {
    int n = graph.size();
    std::vector<int> subtree(n);
    std::vector<bool> removed(n);
    std::vector<int> result;
    result.reserve(n);

    auto calc_subtree = [&](auto&& self, int u, int p) -> int {
        subtree[u] = 1;
        for (const auto& e: graph[u]) {
            if (e.to == p || removed[e.to]) continue;
            subtree[u] += self(self, e.to, u);
        }
        return subtree[u];
    };

    std::vector<int> cur;
    cur.push_back(0);
    while (!cur.empty()) {
        std::vector<int> next;
        for (auto root: cur) {
            calc_subtree(calc_subtree, root, root);
            int centroid = root;
            bool flag = true;
            while (flag) {
                flag = false;
                for (const auto& e: graph[centroid]) {
                    if (removed[e.to] || subtree[e.to] > subtree[centroid]) continue;
                    if (subtree[root] <= 2 * subtree[e.to]) {
                        centroid = e.to;
                        flag = true;
                    }
                }
            }
            result.push_back(centroid);
            removed[centroid] = true;
            for (const auto& e: graph[centroid]) {
                if (removed[e.to]) continue;
                next.push_back(e.to);
            }
        }
        cur = std::move(next);
    }
    return result;
}

}  //  namespace nono

namespace nono {

void solve() {
    int n, k;
    std::cin >> n >> k;
    std::vector<WeightedEdge<int>> edges;
    edges.reserve(n - 1);
    for (int i = 0; i + 1 < n; i++) {
        int u, v, c;
        std::cin >> u >> v >> c;
        u--;
        v--;
        edges.emplace_back(u, v, c);
    }
    const auto graph = to_undirected_graph(n, edges);
    std::vector<bool> removed(n);
    long long ans = 0;
    for (auto centroid: centroids(graph)) {
        std::map<std::pair<int, int>, int> a;

        std::map<int, int> b;

        std::map<int, int> c;
        int d = 0;

        std::vector<std::pair<int, int>> path;
        auto dfs = [&](auto&& self, int u, int p, std::pair<int, int> color) -> void {
            path.push_back(color);
            for (const auto& e: graph[u]) {
                if (e.to == p || removed[e.to]) continue;
                auto new_color = color;
                if (new_color.first == -1) {
                    new_color.first = e.weight;
                    self(self, e.to, u, new_color);
                } else if (new_color.first != e.weight && new_color.second == -1) {
                    new_color.second = e.weight;
                    self(self, e.to, u, new_color);
                } else if (new_color.first == e.weight || new_color.second == e.weight) {
                    self(self, e.to, u, new_color);
                }
            }
        };

        for (const auto& e: graph[centroid]) {
            if (removed[e.to]) continue;
            path.clear();
            dfs(dfs, e.to, centroid, {e.weight, -1});
            for (auto [c1, c2]: path) {
                if (c1 == -1) continue;
                if (c2 == -1) {
                    ans += b[c1];

                    ans += d - c[c1];
                } else {
                    if (c1 > c2) std::swap(c1, c2);

                    ans += a[std::pair(c1, c2)];

                    ans += c[c1];

                    ans += c[c2];

                    ans++;
                }
            }

            for (auto [c1, c2]: path) {
                if (c1 == -1) continue;
                if (c2 == -1) {
                    c[c1]++;
                    d++;
                } else {
                    if (c1 > c2) std::swap(c1, c2);
                    a[std::pair(c1, c2)]++;
                    b[c1]++;
                    b[c2]++;
                }
            }
        }

        removed[centroid] = true;
    }

    std::cout << ans << '\n';
}

}  //  namespace nono

int main() {
    std::cin.tie(0)->sync_with_stdio(0);
    nono::solve();
}
0