CP-Library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub joji0/CP-Library

:heavy_check_mark: test/1_library_checker/graph/tree_diameter.test.cpp

Depends on

Code

// 我思う、故に我在り
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
#include <bits/stdc++.h>
#include "graph/tree/TreeDiameter.hpp"

void solve() {
        int N; std::cin >> N;
        Graph<int64_t> G(N);
        for (int i = 0; i < N - 1; i++) {
                int a, b; int64_t c; std::cin >> a >> b >> c;
                G.add_edge(a, b, c);
        }
        TreeDiameter<int64_t> diam(G);
        std::cout << diam.diameter << " " << diam.path_nodes.size() << '\n';
        for (int v : diam.path_nodes) {
                std::cout << v << " ";
        }
        std::cout << '\n';
}

int main() {
        std::ios::sync_with_stdio(false); std::cin.tie(nullptr);
        int tc = 1;
        // std::cin >> tc;
        while (tc--) solve();
        return 0;
}
#line 1 "test/1_library_checker/graph/tree_diameter.test.cpp"
// 我思う、故に我在り
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
#include <bits/stdc++.h>
#line 5 "graph/Graph.hpp"

template <typename T = int>
struct Edge {
        int from, to;
        T cost;
        int id;
        Edge() : from(-1), to(-1), cost(0), id(-1) {}
        Edge(int _from, int _to, const T& _cost = 1, int _id = -1)
                : from(_from), to(_to), cost(_cost), id(_id) {}
        Edge(int _from, int _to, T&& _cost, int _id = -1)
                : from(_from), to(_to), cost(std::move(_cost)), id(_id) {}
        operator int() const { return to; }
        friend bool operator<(const Edge<T>& lhs, const Edge<T>& rhs) {
                return lhs.cost < rhs.cost;
        }
        friend bool operator>(const Edge<T>& lhs, const Edge<T>& rhs) {
                return lhs.cost > rhs.cost;
        }
};

template <typename T = int>
using Edges  = std::vector<Edge<T>>;

template <typename T = int>
using GMatrix = std::vector<std::vector<T>>;

template <typename T = int>
struct Graph : std::vector<std::vector<Edge<T>>> {
        using Base = std::vector<std::vector<Edge<T>>>;
        using Base::Base;
        int eid = 0;
        int edge_size() const { return eid; }
        int add_edge(int a, int b, const T& c = 1, bool is_directed = false) {
                assert(0 <= a && a < (int)this->size());
                assert(0 <= b && b < (int)this->size());
                (*this)[a].emplace_back(a, b, c, eid);
                if (!is_directed) {
                        (*this)[b].emplace_back(b, a, c, eid);
                }
                return eid++;
        }
};
#line 6 "graph/tree/TreeDiameter.hpp"

template <typename T>
class TreeDiameter {
public:
        T diameter;
        int u, v;
        std::vector<int> path_nodes;
        std::vector<Edge<T>> path_edges;
        TreeDiameter(const Graph<T>& G) {
                int n = G.size();
                if (n == 0) {
                        diameter = T(0);
                        u = v = -1;
                        return;
                }
                if (n == 1) {
                        diameter = T(0);
                        u = v = 0;
                        path_nodes = {0};
                        return;
                }
                auto bfs = [&](int source) {
                        std::vector<T> dist(n, T(0));
                        std::vector<int> parent_node(n, -1);
                        std::vector<Edge<T>> parent_edge(n);
                        std::vector<int> q;
                        q.reserve(n);
                        q.push_back(source);
                        int head = 0;
                        int farthest = source;
                        T max_dist = T(0);
                        while (head < (int)q.size()) {
                                int curr = q[head++];
                                if (dist[curr] >= max_dist) {
                                        max_dist = dist[curr];
                                        farthest = curr;
                                }
                                for (const auto& e : G[curr]) {
                                        int nxt = e.to;
                                        if (nxt == parent_node[curr]) continue;
                                        dist[nxt] = dist[curr] + e.cost;
                                        parent_node[nxt] = curr;
                                        parent_edge[nxt] = e;
                                        q.emplace_back(nxt);
                                }
                        }
                        return std::make_tuple(farthest, max_dist, parent_node, parent_edge, dist);
                };
                auto [u_node, d1, p_node1, p_edge1, dist1] = bfs(0);
                auto [v_node, max, p_node2, p_edge2, dist2] = bfs(u_node);
                diameter = max;
                u = u_node;
                v = v_node;
                int curr = v;
                while (curr != u) {
                        path_nodes.emplace_back(curr);
                        path_edges.emplace_back(p_edge2[curr]);
                        curr = p_node2[curr];
                }
                path_nodes.emplace_back(u);
                std::reverse(path_nodes.begin(), path_nodes.end());
                std::reverse(path_edges.begin(), path_edges.end());
        }
        std::pair<std::vector<int>, std::vector<int>> dist_from_endpoints(const Graph<T>& G) {
                auto dfs = [&](auto& dfs, int v, int p, std::vector<int>& d) -> void {
                        for (int u : G[v]) {
                                if (u == p) continue;
                                d[u] = d[v] + 1;
                                dfs(dfs, u, v, d);
                        }
                };
                std::vector<int> dist_from_u(G.size(), 0), dist_from_v(G.size(), 0);
                dfs(dfs, u, -1, dist_from_u);
                dfs(dfs, v, -1, dist_from_v);
                return {dist_from_u, dist_from_v};
        };
};
#line 5 "test/1_library_checker/graph/tree_diameter.test.cpp"

void solve() {
        int N; std::cin >> N;
        Graph<int64_t> G(N);
        for (int i = 0; i < N - 1; i++) {
                int a, b; int64_t c; std::cin >> a >> b >> c;
                G.add_edge(a, b, c);
        }
        TreeDiameter<int64_t> diam(G);
        std::cout << diam.diameter << " " << diam.path_nodes.size() << '\n';
        for (int v : diam.path_nodes) {
                std::cout << v << " ";
        }
        std::cout << '\n';
}

int main() {
        std::ios::sync_with_stdio(false); std::cin.tie(nullptr);
        int tc = 1;
        // std::cin >> tc;
        while (tc--) solve();
        return 0;
}
Back to top page