4.1.8 Prufer Code

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

/*

Encode and decode labeled trees using Prufer codes. A Prufer code is a sequence of length $n - 2$
that uniquely represents a labeled tree on nodes $0, 1, \ldots, n - 1$. It is useful for counting
labeled trees, generating test cases, and converting between trees and compact sequences. Encoding
repeatedly removes a leaf and records its neighbor until two nodes remain; decoding reverses this,
deducing each removed leaf from the node degrees implied by the code. Both functions choose the
smallest available leaf at each step. This makes the implementation deterministic and matches the
usual textbook convention.

- `encode_prufer()` returns the Prufer code for the global, bidirectionally pre-populated adjacency
  list `adj`, which must form a valid tree and whose indices represent the nodes.
- `decode_prufer()` takes a Prufer code and returns the corresponding tree edges.

The encoded tree must contain at least two nodes, since the empty code represents the unique tree on
two nodes.

Time Complexity:
- O(n log n) per call to `encode_prufer()` or `decode_prufer()`, where $n$ is the number of nodes in
  the tree.

Space Complexity:
- O(n) auxiliary for `encode_prufer()` and `decode_prufer()`.

*/

#include <set>
#include <utility>
#include <vector>

std::vector<int> encode_prufer(const std::vector<std::vector<int>> &adj) {
  int n = static_cast<int>(adj.size());
  std::vector<int> degree(n), parent(n, -1), code;
  std::set<int> leaves;
  if (n <= 2) {
    return code;
  }
  // Root at node n - 1, which is guaranteed to survive smallest-leaf removal (it is never the
  // smallest leaf while other nodes remain). Rooting elsewhere risks recording parent[root] == -2
  // if the chosen root itself gets stripped to a leaf (e.g. when n - 1 happens to be a leaf).
  int root = n - 1;
  for (int u = 0; u < n; u++) {
    degree[u] = static_cast<int>(adj[u].size());
    if (degree[u] == 1) {
      leaves.insert(u);
    }
  }
  std::vector<int> st(1, root);
  parent[root] = -2;
  while (!st.empty()) {
    int u = st.back();
    st.pop_back();
    for (int v : adj[u]) {
      if (parent[v] == -1) {
        parent[v] = u;
        st.push_back(v);
      }
    }
  }
  for (int i = 0; i < n - 2; i++) {
    int leaf = *leaves.begin();
    leaves.erase(leaves.begin());
    int p = parent[leaf];
    code.push_back(p);
    if (--degree[p] == 1) {
      leaves.insert(p);
    }
  }
  return code;
}

std::vector<std::pair<int, int>> decode_prufer(const std::vector<int> &code) {
  int n = static_cast<int>(code.size()) + 2;
  std::vector<int> degree(n, 1);
  std::set<int> leaves;
  std::vector<std::pair<int, int>> edges;
  for (int c : code) {
    degree[c]++;
  }
  for (int u = 0; u < n; u++) {
    if (degree[u] == 1) {
      leaves.insert(u);
    }
  }
  for (int u : code) {
    int leaf = *leaves.begin();
    leaves.erase(leaves.begin());
    edges.emplace_back(leaf, u);
    if (--degree[u] == 1) {
      leaves.insert(u);
    }
  }
  edges.emplace_back(*leaves.begin(), *leaves.rbegin());
  return edges;
}

/*** Example Usage ***/

#include <cassert>
using namespace std;

int main() {
  //     0
  //     |
  // 1---3---2
  //     |
  //     4
  vector<vector<int>> adj(5);
  auto add_edge = [&](int u, int v) {
    adj[u].push_back(v);
    adj[v].push_back(u);
  };
  add_edge(0, 3);
  add_edge(1, 3);
  add_edge(2, 3);
  add_edge(3, 4);
  vector<int> code = encode_prufer(adj);
  assert((code == vector<int>{3, 3, 3}));
  assert((decode_prufer(code) == vector<pair<int, int>>{{0, 3}, {1, 3}, {2, 3}, {3, 4}}));
  return 0;
}