4.4.3 Kruskal Reconstruction Tree

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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

/*

Given a connected, undirected, weighted graph, build the Kruskal reconstruction tree of its minimum
spanning tree. Kruskal's algorithm adds MST edges in nondecreasing weight order; whenever it joins
two components, create a new internal node whose children are the two component roots and whose
value is the joining edge weight. The original graph nodes become leaves.

The lowest common ancestor of two leaves is exactly the merge step that first connects their MST
components, so its value is the maximum edge weight on the MST path between those two nodes. This
turns bottleneck path queries into ordinary LCA queries on the reconstruction tree.

- `KruskalReconstructionTree(n, edges)` builds the tree from a connected graph whose weighted edges
  are stored as (`weight`, `u`, `v`), where the `n` original nodes are numbered $[0, `n`)$. Parallel
  edges are supported.
- `max_edge_on_path(u, v)` returns the maximum edge weight on the MST path between original nodes
  `u` and `v`, which must be distinct.

Time Complexity:
- O(m log m + n log n) for construction, where $n$ is the number of nodes and $m$ is the number of
  edges.
- O(log n) per call to `max_edge_on_path()`.

Space Complexity:
- O(n log n) object storage for the LCA table and node values, plus O(n + m) auxiliary during
  construction.

*/

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <numeric>
#include <tuple>
#include <vector>

class KruskalReconstructionTree {
  std::vector<int> tin, tout;
  std::vector<std::vector<int>> up;
  std::vector<int64_t> value;

  bool is_ancestor(int u, int v) const { return tin[u] <= tin[v] && tout[v] <= tout[u]; }

  int lca(int u, int v) const {
    if (is_ancestor(u, v)) {
      return u;
    }
    if (is_ancestor(v, u)) {
      return v;
    }
    for (int k = static_cast<int>(up.size()) - 1; k >= 0; k--) {
      if (!is_ancestor(up[k][u], v)) {
        u = up[k][u];
      }
    }
    return up[0][u];
  }

 public:
  KruskalReconstructionTree(int n, std::vector<std::tuple<int64_t, int, int>> edges) {
    assert(n > 0);
    int total_nodes = 2 * n - 1;
    std::vector<int> dsu_root(total_nodes), dsu_tree_root(total_nodes);
    std::vector<std::vector<int>> tree(total_nodes);
    value.resize(total_nodes);
    std::iota(dsu_root.begin(), dsu_root.begin() + n, 0);
    std::iota(dsu_tree_root.begin(), dsu_tree_root.begin() + n, 0);
    auto find = [&](auto &&find, int u) -> int {
      return dsu_root[u] == u ? u : dsu_root[u] = find(find, dsu_root[u]);
    };
    std::sort(edges.begin(), edges.end());
    int nodes = n;
    for (auto [w, u, v] : edges) {
      u = find(find, u);
      v = find(find, v);
      if (u == v) {
        continue;
      }
      value[nodes] = w;
      tree[nodes].push_back(dsu_tree_root[u]);
      tree[nodes].push_back(dsu_tree_root[v]);
      dsu_root[u] = nodes;
      dsu_root[v] = nodes;
      dsu_root[nodes] = nodes;
      dsu_tree_root[nodes] = nodes;
      nodes++;
    }
    assert(nodes == total_nodes);
    int lg = 1;
    while ((1 << lg) <= nodes) {
      lg++;
    }
    tin.resize(nodes);
    tout.resize(nodes);
    up.assign(lg, std::vector<int>(nodes));
    int timer = 0;
    auto dfs = [&](auto &&dfs, int u, int p) -> void {
      tin[u] = timer++;
      up[0][u] = p;
      for (int k = 1; k < lg; k++) {
        up[k][u] = up[k - 1][up[k - 1][u]];
      }
      for (int v : tree[u]) {
        dfs(dfs, v, u);
      }
      tout[u] = timer;
    };
    dfs(dfs, nodes - 1, nodes - 1);
  }

  int64_t max_edge_on_path(int u, int v) const {
    assert(u != v);
    return value[lca(u, v)];
  }
};

/*** Example Usage ***/

using namespace std;

int main() {
  //            0
  //    w=2   /  | w=8
  //        /    |
  //      1------2
  //      | w=4 /
  //  w=5 |   /
  //      | / w=9
  //      3
  vector<tuple<int64_t, int, int>> edges{{2, 0, 1}, {4, 1, 2}, {5, 1, 3}, {8, 0, 2}, {9, 2, 3}};
  KruskalReconstructionTree tree(4, edges);

  // Reconstruction tree; internal nodes are labeled with their joining edge weight.
  //              6 (5)
  //             / |
  //        5 (4)  3
  //       / |
  //  4 (2)  2
  //  / |
  // 0  1
  assert(tree.max_edge_on_path(0, 2) == 4);
  assert(tree.max_edge_on_path(0, 3) == 5);
  assert(tree.max_edge_on_path(2, 3) == 5);
  return 0;
}