4.1.4 Tree Centers and Diameter

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

/*

An unweighted tree possesses a center, centroid, and diameter. The centers are found by repeatedly
peeling away the leaves of the tree until one or two nodes remain. The diameter relies on the fact
that the farthest node from any start is always one endpoint of a diameter, so a second traversal
from that endpoint measures the full diameter. The following functions apply to a global,
bidirectionally pre-populated adjacency list `adj` which must form a valid tree with nodes numbered
$[0, n)$, where $n$ is `adj.size()`.

- `find_centers()` returns a vector of either one or two tree Jordan centers. The Jordan center of a
  tree is the set of all nodes with minimum eccentricity, that is, the set of all nodes where the
  maximum distance to all other nodes in the tree is minimal.
- `find_centroid()` returns the node where all of its subtrees have a size less than or equal to
  $n / 2$, where $n$ is the number of nodes in the tree.
- `diameter()` returns a tuple (`length`, `u`, `v`), where `length` is the maximum distance between
  any two nodes in the tree and `u` and `v` are endpoints attaining that distance. It uses a
  well-known double DFS technique.

Time Complexity:
- O(max(n, m)) per call to `find_centers()`, `find_centroid()`, and `diameter()`, where $n$ is the
  number of nodes and $m$ is the number of edges.

Space Complexity:
- O(n) auxiliary heap space for `find_centers()`, where $n$ is the number of nodes.
- O(n) auxiliary stack space for `find_centroid()` and `diameter()`.

*/

#include <algorithm>
#include <tuple>
#include <utility>
#include <vector>

std::vector<std::vector<int>> adj;

std::vector<int> find_centers() {
  int n = static_cast<int>(adj.size());
  std::vector<int> leaves, degree(n);
  for (int i = 0; i < n; i++) {
    degree[i] = static_cast<int>(adj[i].size());
    if (degree[i] <= 1) {
      leaves.push_back(i);
    }
  }
  int removed = static_cast<int>(leaves.size());
  while (removed < n) {
    std::vector<int> nleaves;
    for (int u : leaves) {
      for (int v : adj[u]) {
        if (--degree[v] == 1) {
          nleaves.push_back(v);
        }
      }
    }
    leaves = nleaves;
    removed += static_cast<int>(leaves.size());
  }
  return leaves;
}

// Returns the centroid node index if found in this subtree, or -(subtree size) to propagate the
// size up to the parent so it can check the complementary component's size.
int find_centroid(int u = 0, int p = -1) {
  int n = static_cast<int>(adj.size());
  int count = 1;
  bool good_center = true;
  for (int v : adj[u]) {
    if (v == p) {
      continue;
    }
    int res = find_centroid(v, u);
    if (res >= 0) {
      return res;
    }
    int size = -res;
    good_center &= (size <= n / 2);
    count += size;
  }
  good_center &= (n - count <= n / 2);
  return good_center ? u : -count;
}

std::tuple<int, int, int> diameter() {
  auto dfs = [&](auto &&dfs, int u, int p = -1, int depth = 0) -> std::pair<int, int> {
    std::pair<int, int> res{depth, u};
    for (int v : adj[u]) {
      if (v != p) {
        res = std::max(res, dfs(dfs, v, u, depth + 1));
      }
    }
    return res;
  };
  int u = dfs(dfs, 0).second;
  auto [length, v] = dfs(dfs, u);
  return {length, u, v};
}

/*** Example Usage ***/

#include <cassert>
using namespace std;

void add_edge(int u, int v) {
  adj[u].push_back(v);
  adj[v].push_back(u);
}

int main() {
  // 0---1---4---3
  //     |   |
  //     2   5
  adj.assign(6, {});
  add_edge(0, 1);
  add_edge(1, 2);
  add_edge(1, 4);
  add_edge(3, 4);
  add_edge(4, 5);
  assert((find_centers() == vector<int>{1, 4}));
  assert(find_centroid() == 4);
  auto [length, u, v] = diameter();
  assert(length == 3 && u == 5 && v == 2);
  return 0;
}