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
145
146
147
148

/*

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 weighted edges stored as (`weight`,
  `u`, `v`).
- `root()` returns the reconstruction-tree root.
- `max_edge_on_path(u, v)` returns the maximum edge weight on the MST path between original nodes
  `u` and `v`.

The input graph must be connected. Parallel edges are supported and are handled naturally by
Kruskal's algorithm.

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) for the reconstruction tree and LCA table, plus O(m) auxiliary during construction.

*/

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

class KruskalReconstructionTree {
  std::vector<int> dsu_root, dsu_tree_root, depth, tin, tout;
  std::vector<std::vector<int>> up, tree;
  std::vector<int64_t> value;
  int timer = 0, root_node = -1;

  int find(int u) { return dsu_root[u] == u ? u : dsu_root[u] = find(dsu_root[u]); }

  void dfs_lca(int u, int p) {
    tin[u] = timer++;
    up[0][u] = p;
    for (int k = 1; k < static_cast<int>(up.size()); k++) {
      up[k][u] = up[k - 1][up[k - 1][u]];
    }
    for (int v : tree[u]) {
      depth[v] = depth[u] + 1;
      dfs_lca(v, u);
    }
    tout[u] = timer;
  }

  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)
      : dsu_root(2 * n - 1),
        dsu_tree_root(2 * n - 1),
        depth(2 * n - 1),
        tin(2 * n - 1),
        tout(2 * n - 1),
        tree(2 * n - 1),
        value(2 * n - 1) {
    assert(n > 0);
    std::iota(dsu_root.begin(), dsu_root.begin() + n, 0);
    std::iota(dsu_tree_root.begin(), dsu_tree_root.begin() + n, 0);
    std::sort(edges.begin(), edges.end());
    int nodes = n;
    for (auto [w, u, v] : edges) {
      u = find(u);
      v = 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 == 2 * n - 1);
    root_node = nodes - 1;
    int lg = 1;
    while ((1 << lg) <= nodes) {
      lg++;
    }
    up.assign(lg, std::vector<int>(nodes));
    dfs_lca(root_node, root_node);
  }

  int root() const { return root_node; }
  int64_t max_edge_on_path(int u, int v) const { 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.root() == 6);
  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;
}