Tree Diameter
(src/tree/TreeDiameter.hpp)

• View this file on GitHub
• Last update: 2023-01-12 22:28:24+09:00
• Include: #include "src/tree/TreeDiameter.hpp"

概要

木の直径を計算するライブラリ．

メンバ関数	効果	時間計算量
TreeDiameter(n)	頂点数 $n$ の木として初期化する．	$\mathrm{O}(n)$
add_edge(u, b)	辺 $(u,v)$ を追加する．	$\mathrm{O}(1)$
get_diameter_path()	木の直径の長さとそのパスをペアにして返す．	$\mathrm{O}(n)$
farthest_distance()	各頂点から最も遠い頂点までの距離を返す．この最も遠い頂点は直径を構成する端点のどちらかであることが直径の最長性から示せる．	$\mathrm{O}(n)$

出題例

• Codeforces Round #292 (Div. 1) D. Drazil and Morning Exercise
• AtCoder Regular Contest 108 F - Paint Tree

Verified with

• test/aoj/GRL_5_B.test.cpp
• test/yosupo/tree_diameter.test.cpp

Code

#pragma once
#include <cassert>
#include <vector>

template <typename T = int> struct TreeDiameter {
    std::vector<std::vector<std::pair<int, T>>> G;

    TreeDiameter(int n) : G(n), n(n), dist(n), par(n) {}

    void add_edge(int u, int v, T c = 1) {
        assert(0 <= u && u < n);
        assert(0 <= v && v < n);
        G[u].emplace_back(v, c);
        G[v].emplace_back(u, c);
    }

    std::pair<T, std::vector<int>> get_diameter_path() {
        argmax = 0;
        dfs(argmax, -1);
        dfs(argmax, -1);
        T diam = dist[argmax];
        std::vector<int> path;
        while (argmax >= 0) {
            path.emplace_back(argmax);
            argmax = par[argmax];
        }
        return {diam, path};
    }

    std::vector<T> farthest_distance() {
        auto path = get_diameter_path().second;
        int s = path.front(), t = path.back();
        dfs(s, -1);
        auto dist_s = dist;
        dfs(t, -1);
        auto dist_t = dist;
        for (int i = 0; i < n; i++) dist_s[i] = std::max(dist_s[i], dist_t[i]);
        return dist_s;
    }

private:
    int n, argmax;
    std::vector<T> dist;
    std::vector<int> par;

    void dfs(int v, int p) {
        par[v] = p;
        if (p < 0) dist[v] = T(0);
        if (dist[argmax] < dist[v]) argmax = v;
        for (auto& e : G[v]) {
            int u = e.first;
            if (u == p) continue;
            dist[u] = dist[v] + e.second;
            dfs(u, v);
        }
    }
};

#line 2 "src/tree/TreeDiameter.hpp"
#include <cassert>
#include <vector>

template <typename T = int> struct TreeDiameter {
    std::vector<std::vector<std::pair<int, T>>> G;

    TreeDiameter(int n) : G(n), n(n), dist(n), par(n) {}

    void add_edge(int u, int v, T c = 1) {
        assert(0 <= u && u < n);
        assert(0 <= v && v < n);
        G[u].emplace_back(v, c);
        G[v].emplace_back(u, c);
    }

    std::pair<T, std::vector<int>> get_diameter_path() {
        argmax = 0;
        dfs(argmax, -1);
        dfs(argmax, -1);
        T diam = dist[argmax];
        std::vector<int> path;
        while (argmax >= 0) {
            path.emplace_back(argmax);
            argmax = par[argmax];
        }
        return {diam, path};
    }

    std::vector<T> farthest_distance() {
        auto path = get_diameter_path().second;
        int s = path.front(), t = path.back();
        dfs(s, -1);
        auto dist_s = dist;
        dfs(t, -1);
        auto dist_t = dist;
        for (int i = 0; i < n; i++) dist_s[i] = std::max(dist_s[i], dist_t[i]);
        return dist_s;
    }

private:
    int n, argmax;
    std::vector<T> dist;
    std::vector<int> par;

    void dfs(int v, int p) {
        par[v] = p;
        if (p < 0) dist[v] = T(0);
        if (dist[argmax] < dist[v]) argmax = v;
        for (auto& e : G[v]) {
            int u = e.first;
            if (u == p) continue;
            dist[u] = dist[v] + e.second;
            dfs(u, v);
        }
    }
};

Back to top page