prufer_code.hpp

SECTIONGraph INCLUDEnoya/prufer_code.hpp

Encode a labeled tree on vertices [0, n) as its Prüfer sequence in O(n log n), always removing the smallest current leaf.

在有标号树与 Prüfer 序列之间互相转换；用于树计数、按度数构造或枚举。

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_PRUFER_CODE_HPP
#define NOYA_PRUFER_CODE_HPP 1

/// @complexity Time: O(n log n).
/// Space: O(n).

#include <cassert>
#include <functional>
#include <queue>
#include <utility>
#include <vector>

namespace noya {

/// @brief Encode a labeled tree on vertices [0, n) as its Prüfer sequence in
/// O(n log n), always removing the smallest current leaf.
inline std::vector<int>
prufer_encode(int n, const std::vector<std::pair<int, int>> &edges) {
  assert(n >= 2);
  assert(int(edges.size()) == n - 1);
  if (n <= 2) {
    for (auto [first, second] : edges) {
      assert(0 <= first && first < n);
      assert(0 <= second && second < n);
      assert(first != second);
    }
    return {};
  }
  std::vector<std::vector<int>> graph(n);
  std::vector<int> degree(n);
  for (auto [first, second] : edges) {
    assert(0 <= first && first < n);
    assert(0 <= second && second < n);
    assert(first != second);
    graph[first].push_back(second);
    graph[second].push_back(first);
    degree[first]++;
    degree[second]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> leaves;
  for (int vertex = 0; vertex < n; vertex++) {
    if (degree[vertex] == 1) {
      leaves.push(vertex);
    }
  }
  std::vector<int> code;
  code.reserve(n - 2);
  for (int step = 0; step < n - 2; step++) {
    assert(!leaves.empty());
    int leaf = leaves.top();
    leaves.pop();
    int neighbor = -1;
    for (int next : graph[leaf]) {
      if (degree[next] > 0) {
        neighbor = next;
        break;
      }
    }
    assert(neighbor != -1);
    code.push_back(neighbor);
    degree[leaf] = 0;
    if (--degree[neighbor] == 1) {
      leaves.push(neighbor);
    }
  }
  return code;
}

/// @brief Decode a Prüfer sequence over labels [0, code.size()+2) into a tree
/// in O(n log n), returning edges in leaf-removal order.
inline std::vector<std::pair<int, int>>
prufer_decode(const std::vector<int> &code) {
  const int n = int(code.size()) + 2;
  std::vector<int> degree(n, 1);
  for (int vertex : code) {
    assert(0 <= vertex && vertex < n);
    degree[vertex]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> leaves;
  for (int vertex = 0; vertex < n; vertex++) {
    if (degree[vertex] == 1) {
      leaves.push(vertex);
    }
  }
  std::vector<std::pair<int, int>> edges;
  edges.reserve(n - 1);
  for (int vertex : code) {
    int leaf = leaves.top();
    leaves.pop();
    edges.emplace_back(leaf, vertex);
    degree[leaf]--;
    if (--degree[vertex] == 1) {
      leaves.push(vertex);
    }
  }
  int first = leaves.top();
  leaves.pop();
  int second = leaves.top();
  edges.emplace_back(first, second);
  return edges;
}

} // namespace noya

#endif // NOYA_PRUFER_CODE_HPP

#include <cassert>
#include <functional>
#include <queue>
#include <utility>
#include <vector>

/// @complexity Time: O(n log n).
/// Space: O(n).

namespace noya {

/// @brief Encode a labeled tree on vertices [0, n) as its Prüfer sequence in
/// O(n log n), always removing the smallest current leaf.
inline std::vector<int>
prufer_encode(int n, const std::vector<std::pair<int, int>> &edges) {
  assert(n >= 2);
  assert(int(edges.size()) == n - 1);
  if (n <= 2) {
    for (auto [first, second] : edges) {
      assert(0 <= first && first < n);
      assert(0 <= second && second < n);
      assert(first != second);
    }
    return {};
  }
  std::vector<std::vector<int>> graph(n);
  std::vector<int> degree(n);
  for (auto [first, second] : edges) {
    assert(0 <= first && first < n);
    assert(0 <= second && second < n);
    assert(first != second);
    graph[first].push_back(second);
    graph[second].push_back(first);
    degree[first]++;
    degree[second]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> leaves;
  for (int vertex = 0; vertex < n; vertex++) {
    if (degree[vertex] == 1) {
      leaves.push(vertex);
    }
  }
  std::vector<int> code;
  code.reserve(n - 2);
  for (int step = 0; step < n - 2; step++) {
    assert(!leaves.empty());
    int leaf = leaves.top();
    leaves.pop();
    int neighbor = -1;
    for (int next : graph[leaf]) {
      if (degree[next] > 0) {
        neighbor = next;
        break;
      }
    }
    assert(neighbor != -1);
    code.push_back(neighbor);
    degree[leaf] = 0;
    if (--degree[neighbor] == 1) {
      leaves.push(neighbor);
    }
  }
  return code;
}

/// @brief Decode a Prüfer sequence over labels [0, code.size()+2) into a tree
/// in O(n log n), returning edges in leaf-removal order.
inline std::vector<std::pair<int, int>>
prufer_decode(const std::vector<int> &code) {
  const int n = int(code.size()) + 2;
  std::vector<int> degree(n, 1);
  for (int vertex : code) {
    assert(0 <= vertex && vertex < n);
    degree[vertex]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> leaves;
  for (int vertex = 0; vertex < n; vertex++) {
    if (degree[vertex] == 1) {
      leaves.push(vertex);
    }
  }
  std::vector<std::pair<int, int>> edges;
  edges.reserve(n - 1);
  for (int vertex : code) {
    int leaf = leaves.top();
    leaves.pop();
    edges.emplace_back(leaf, vertex);
    degree[leaf]--;
    if (--degree[vertex] == 1) {
      leaves.push(vertex);
    }
  }
  int first = leaves.top();
  leaves.pop();
  int second = leaves.top();
  edges.emplace_back(first, second);
  return edges;
}

} // namespace noya