Tree Diameter (#39)

* Tree Diameter

* [auto-verifier] verify commit 7231abe

Co-authored-by: GitHub <noreply@github.com>

Author: hitonanode

.verify-helper/timestamps.remote.json

@@ -133,6 +133,7 @@
 "string/test/sa_count_keyword.test.cpp": "2021-01-02 01:50:58 +0900",
 "string/test/suffix_array.test.cpp": "2021-01-02 00:51:41 +0900",
 "string/test/z_algorithm.test.cpp": "2021-01-01 16:52:32 +0900",
+"tree/test/diameter.test.cpp": "2021-06-09 00:24:00 +0900",
 "tree/test/frequency_table_of_tree_distance.test.cpp": "2021-06-06 15:23:40 +0900",
 "tree/test/hl_decomposition.test.cpp": "2021-06-06 15:23:40 +0900",
 "tree/test/lca.test.cpp": "2021-02-26 00:36:27 +0900",

tree/diameter.hpp

@@ -0,0 +1,50 @@
+#pragma once
+#include <algorithm>
+#include <cassert>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+// CUT begin
+// 木の直径と，これをなすパスの頂点を出力
+template <class Weight> struct TreeDiameter {
+    int n;
+    std::vector<std::vector<std::pair<int, Weight>>> to;
+    std::vector<int> prev;
+    std::vector<Weight> dist;
+    TreeDiameter(int V) : n(V), to(V), prev(V), dist(V) {}
+    void add_edge(int s, int t, Weight w) {
+        assert(0 <= s and s < n);
+        assert(0 <= t and t < n);
+        to[s].emplace_back(t, w);
+        to[t].emplace_back(s, w);
+    }
+    int argdmax;
+    Weight dmax;
+    void dfs(int now, int prv) {
+        if (dmax < dist[now]) dmax = dist[now], argdmax = now;
+        for (auto p : to[now]) {
+            int nxt;
+            Weight w;
+            std::tie(nxt, w) = p;
+            if (nxt == prv) continue;
+            prev[nxt] = now;
+            dist[nxt] = dist[now] + w;
+            dfs(nxt, now);
+        }
+    }
+    std::pair<Weight, std::vector<int>> get_diameter_edges(int root = 0) {
+        prev[root] = -1;
+        dist[root] = 0;
+        dmax = 0, argdmax = root;
+        dfs(root, -1);
+        dmax = 0, prev[argdmax] = -1, dist[argdmax] = 0;
+        dfs(argdmax, -1);
+        std::vector<int> ret;
+        while (argdmax >= 0) {
+            ret.push_back(argdmax);
+            argdmax = prev[argdmax];
+        }
+        return {dmax, ret};
+    }
+};

tree/diameter.md

@@ -0,0 +1,20 @@
+---
+title: Tree diameter （木の直径）
+documentation_of: ./diameter.hpp
+---
+
+木の直径を計算．
+
+## 使用方法
+
+```cpp
+TreeDiameter<long long> tree(N);
+for (auto [s, t, w] : edges) tree.add_edge(s, t, w);
+
+// 頂点 r を含む部分木に関して，直径 D とこれをなす頂点列 vs （の一つ）を計算
+auto [D, vs] = tree.get_diameter_edges(r);
+```
+
+## 問題例
+
+- [Library Checker: Tree Diameter](https://judge.yosupo.jp/problem/tree_diameter)

tree/test/diameter.test.cpp

@@ -0,0 +1,20 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
+#include "../diameter.hpp"
+#include <iostream>
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+    int N;
+    cin >> N;
+    TreeDiameter<long long> tree(N);
+    for (int i = 0; i < N - 1; i++) {
+        int a, b, c;
+        cin >> a >> b >> c;
+        tree.add_edge(a, b, c);
+    }
+    auto ret = tree.get_diameter_edges();
+    cout << ret.first << ' ' << ret.second.size() << '\n';
+    for (auto x : ret.second) cout << x << ' ';
+    cout << '\n';
+}