結果

問題 No.1002 Twotone
ユーザー nono00nono00
提出日時 2024-02-27 03:56:37
言語 C++23
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 904 ms / 5,000 ms
コード長 7,324 bytes
コンパイル時間 2,055 ms
コンパイル使用メモリ 147,980 KB
実行使用メモリ 33,780 KB
最終ジャッジ日時 2024-02-27 03:56:55
合計ジャッジ時間 14,643 ms
ジャッジサーバーID
（参考情報） 		judge14 / judge10
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
6,676 KB
testcase_01 AC 2 ms
6,676 KB
testcase_02 AC 2 ms
6,676 KB
testcase_03 AC 218 ms
11,212 KB
testcase_04 AC 285 ms
13,540 KB
testcase_05 AC 316 ms
13,668 KB
testcase_06 AC 2 ms
6,676 KB
testcase_07 AC 141 ms
9,844 KB
testcase_08 AC 257 ms
14,308 KB
testcase_09 AC 249 ms
14,308 KB
testcase_10 AC 2 ms
6,676 KB
testcase_11 AC 284 ms
11,380 KB
testcase_12 AC 423 ms
13,924 KB
testcase_13 AC 394 ms
13,924 KB
testcase_14 AC 2 ms
6,676 KB
testcase_15 AC 198 ms
9,448 KB
testcase_16 AC 424 ms
13,924 KB
testcase_17 AC 400 ms
13,924 KB
testcase_18 AC 2 ms
6,676 KB
testcase_19 AC 416 ms
16,320 KB
testcase_20 AC 522 ms
21,732 KB
testcase_21 AC 494 ms
21,348 KB
testcase_22 AC 2 ms
6,676 KB
testcase_23 AC 365 ms
16,888 KB
testcase_24 AC 755 ms
24,900 KB
testcase_25 AC 709 ms
24,952 KB
testcase_26 AC 2 ms
6,676 KB
testcase_27 AC 57 ms
10,512 KB
testcase_28 AC 93 ms
14,436 KB
testcase_29 AC 90 ms
14,436 KB
testcase_30 AC 2 ms
6,676 KB
testcase_31 AC 82 ms
14,404 KB
testcase_32 AC 88 ms
14,436 KB
testcase_33 AC 83 ms
14,436 KB
testcase_34 AC 904 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