tree/tree_diameter.hpp

• View this file on GitHub
• Last update: 2026-08-12 07:38:15+00:00
• Include: #include "tree/tree_diameter.hpp"

Depends on

• graph/csr_graph.hpp

Verified with

• test/tree/tree_diameter.test.cpp

Code

#ifndef TREE_DIAMETER_HPP
#define TREE_DIAMETER_HPP

#include "graph/csr_graph.hpp"
#include <algorithm>
#include <cassert>
#include <vector>

// requires non-negative edge weights
// return (diameter, path)
// for an unweighted tree, centers are path[d / 2] and path[(d + 1) / 2]
template <typename EdgeWeight> auto tree_diameter(const CSRGraph<EdgeWeight> &g) {
    assert(0 < g.size());
    auto n = int(g.size());
    std::vector<int> parent(n, -1);
    std::vector<EdgeWeight> depth(n);
    auto dfs = [&](int x) {
        std::vector<int> st;
        st.push_back(x);
        parent[x] = -1;
        depth[x] = 0;
        while (!st.empty()) {
            auto u = st.back();
            st.pop_back();
            for (const auto &e : g[u]) {
                if (e.v != parent[u]) {
                    st.push_back(e.v);
                    parent[e.v] = u;
                    depth[e.v] = depth[u] + e.w;
                }
            }
        }
        auto it = std::ranges::max_element(depth);
        auto r = int(it - depth.begin());
        return std::pair(r, *it);
    };
    auto x = dfs(0).first;
    auto [y, d] = dfs(x);
    std::vector<int> path;
    for (auto u = y; ~u; u = parent[u]) {
        path.push_back(u);
    }
    return std::pair(d, path);
}

#endif // TREE_DIAMETER_HPP

#line 1 "tree/tree_diameter.hpp"



#line 1 "graph/csr_graph.hpp"



#include <cassert>
#include <ranges>
#include <vector>

template <typename EdgeWeight = int> struct CSRGraph {
    struct Edge {
        int u;
        int v;
        EdgeWeight w;
        int i;
    };
    CSRGraph(int n) : n_(n) {}
    void add_edge(int u, int v, EdgeWeight w = 1) {
        assert(!built_);
        assert(0 <= u && u < n_ && 0 <= v && v < n_);
        auto i = int(edges_.size());
        edges_.push_back({u, v, w, i});
    }
    void build_undirected() {
        assert(!built_);
        start_.assign(n_ + 1, 0);
        csr_.resize(2 * edges_.size());
        for (const auto &e : edges_) {
            ++start_[e.u + 1];
            ++start_[e.v + 1];
        }
        for (int i = 0; i < n_; ++i) {
            start_[i + 1] += start_[i];
        }
        auto pos = start_;
        for (const auto &e : edges_) {
            csr_[pos[e.u]++] = {e.u, e.v, e.w, e.i};
            csr_[pos[e.v]++] = {e.v, e.u, e.w, e.i};
        }
        built_ = true;
    }
    void build_directed() {
        assert(!built_);
        start_.assign(n_ + 1, 0);
        csr_.resize(edges_.size());
        for (const auto &e : edges_) {
            ++start_[e.u + 1];
        }
        for (int i = 0; i < n_; ++i) {
            start_[i + 1] += start_[i];
        }
        auto pos = start_;
        for (const auto &e : edges_) {
            csr_[pos[e.u]++] = e;
        }
        built_ = true;
    }
    auto operator[](int u) const {
        assert(built_);
        assert(0 <= u && u < n_);
        return std::ranges::subrange(csr_.begin() + start_[u], csr_.begin() + start_[u + 1]);
    }
    int size() const { return n_; }

private:
    int n_;
    bool built_ = false;
    std::vector<Edge> edges_;
    std::vector<Edge> csr_;
    std::vector<int> start_;
};


#line 5 "tree/tree_diameter.hpp"
#include <algorithm>
#line 8 "tree/tree_diameter.hpp"

// requires non-negative edge weights
// return (diameter, path)
// for an unweighted tree, centers are path[d / 2] and path[(d + 1) / 2]
template <typename EdgeWeight> auto tree_diameter(const CSRGraph<EdgeWeight> &g) {
    assert(0 < g.size());
    auto n = int(g.size());
    std::vector<int> parent(n, -1);
    std::vector<EdgeWeight> depth(n);
    auto dfs = [&](int x) {
        std::vector<int> st;
        st.push_back(x);
        parent[x] = -1;
        depth[x] = 0;
        while (!st.empty()) {
            auto u = st.back();
            st.pop_back();
            for (const auto &e : g[u]) {
                if (e.v != parent[u]) {
                    st.push_back(e.v);
                    parent[e.v] = u;
                    depth[e.v] = depth[u] + e.w;
                }
            }
        }
        auto it = std::ranges::max_element(depth);
        auto r = int(it - depth.begin());
        return std::pair(r, *it);
    };
    auto x = dfs(0).first;
    auto [y, d] = dfs(x);
    std::vector<int> path;
    for (auto u = y; ~u; u = parent[u]) {
        path.push_back(u);
    }
    return std::pair(d, path);
}

Back to top page