prufer_code.hpp

SECTIONGraph INCLUDEnoya/prufer_code.hpp

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

Complexity: Time: O(n log n). Space: O(n).

跳到代码 · GitHub ↗

Implementation

ReferenceSourceExpanded

当前头文件，省略 include guard；依赖见 #include。

/// @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>> &es) {
  assert(n >= 2);
  assert(int(es.size()) == n - 1);
  if (n <= 2) {
    for (auto [a, b] : es) {
      assert(0 <= a && a < n);
      assert(0 <= b && b < n);
      assert(a != b);
    }
    return {};
  }
  std::vector<std::vector<int>> g(n);
  std::vector<int> deg(n);
  for (auto [a, b] : es) {
    assert(0 <= a && a < n);
    assert(0 <= b && b < n);
    assert(a != b);
    g[a].push_back(b);
    g[b].push_back(a);
    deg[a]++;
    deg[b]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> ls;
  for (int u = 0; u < n; u++) {
    if (deg[u] == 1) {
      ls.push(u);
    }
  }
  std::vector<int> s;
  s.reserve(n - 2);
  for (int stp = 0; stp < n - 2; stp++) {
    assert(!ls.empty());
    int lf = ls.top();
    ls.pop();
    int v = -1;
    for (int nxt : g[lf]) {
      if (deg[nxt] > 0) {
        v = nxt;
        break;
      }
    }
    assert(v != -1);
    s.push_back(v);
    deg[lf] = 0;
    if (--deg[v] == 1) {
      ls.push(v);
    }
  }
  return s;
}

/// @brief Decode a Prüfer sequence over labels [0, s.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> &s) {
  const int n = int(s.size()) + 2;
  std::vector<int> deg(n, 1);
  for (int u : s) {
    assert(0 <= u && u < n);
    deg[u]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> ls;
  for (int u = 0; u < n; u++) {
    if (deg[u] == 1) {
      ls.push(u);
    }
  }
  std::vector<std::pair<int, int>> es;
  es.reserve(n - 1);
  for (int u : s) {
    int lf = ls.top();
    ls.pop();
    es.emplace_back(lf, u);
    deg[lf]--;
    if (--deg[u] == 1) {
      ls.push(u);
    }
  }
  int a = ls.top();
  ls.pop();
  int b = ls.top();
  es.emplace_back(a, b);
  return es;
}

} // namespace noya

#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>> &es) {
  assert(n >= 2);
  assert(int(es.size()) == n - 1);
  if (n <= 2) {
    for (auto [a, b] : es) {
      assert(0 <= a && a < n);
      assert(0 <= b && b < n);
      assert(a != b);
    }
    return {};
  }
  std::vector<std::vector<int>> g(n);
  std::vector<int> deg(n);
  for (auto [a, b] : es) {
    assert(0 <= a && a < n);
    assert(0 <= b && b < n);
    assert(a != b);
    g[a].push_back(b);
    g[b].push_back(a);
    deg[a]++;
    deg[b]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> ls;
  for (int u = 0; u < n; u++) {
    if (deg[u] == 1) {
      ls.push(u);
    }
  }
  std::vector<int> s;
  s.reserve(n - 2);
  for (int stp = 0; stp < n - 2; stp++) {
    assert(!ls.empty());
    int lf = ls.top();
    ls.pop();
    int v = -1;
    for (int nxt : g[lf]) {
      if (deg[nxt] > 0) {
        v = nxt;
        break;
      }
    }
    assert(v != -1);
    s.push_back(v);
    deg[lf] = 0;
    if (--deg[v] == 1) {
      ls.push(v);
    }
  }
  return s;
}

/// @brief Decode a Prüfer sequence over labels [0, s.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> &s) {
  const int n = int(s.size()) + 2;
  std::vector<int> deg(n, 1);
  for (int u : s) {
    assert(0 <= u && u < n);
    deg[u]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> ls;
  for (int u = 0; u < n; u++) {
    if (deg[u] == 1) {
      ls.push(u);
    }
  }
  std::vector<std::pair<int, int>> es;
  es.reserve(n - 1);
  for (int u : s) {
    int lf = ls.top();
    ls.pop();
    es.emplace_back(lf, u);
    deg[lf]--;
    if (--deg[u] == 1) {
      ls.push(u);
    }
  }
  int a = ls.top();
  ls.pop();
  int b = ls.top();
  es.emplace_back(a, b);
  return es;
}

} // 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>> &es) {
  assert(n >= 2);
  assert(int(es.size()) == n - 1);
  if (n <= 2) {
    for (auto [a, b] : es) {
      assert(0 <= a && a < n);
      assert(0 <= b && b < n);
      assert(a != b);
    }
    return {};
  }
  std::vector<std::vector<int>> g(n);
  std::vector<int> deg(n);
  for (auto [a, b] : es) {
    assert(0 <= a && a < n);
    assert(0 <= b && b < n);
    assert(a != b);
    g[a].push_back(b);
    g[b].push_back(a);
    deg[a]++;
    deg[b]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> ls;
  for (int u = 0; u < n; u++) {
    if (deg[u] == 1) {
      ls.push(u);
    }
  }
  std::vector<int> s;
  s.reserve(n - 2);
  for (int stp = 0; stp < n - 2; stp++) {
    assert(!ls.empty());
    int lf = ls.top();
    ls.pop();
    int v = -1;
    for (int nxt : g[lf]) {
      if (deg[nxt] > 0) {
        v = nxt;
        break;
      }
    }
    assert(v != -1);
    s.push_back(v);
    deg[lf] = 0;
    if (--deg[v] == 1) {
      ls.push(v);
    }
  }
  return s;
}

/// @brief Decode a Prüfer sequence over labels [0, s.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> &s) {
  const int n = int(s.size()) + 2;
  std::vector<int> deg(n, 1);
  for (int u : s) {
    assert(0 <= u && u < n);
    deg[u]++;
  }
  std::priority_queue<int, std::vector<int>, std::greater<>> ls;
  for (int u = 0; u < n; u++) {
    if (deg[u] == 1) {
      ls.push(u);
    }
  }
  std::vector<std::pair<int, int>> es;
  es.reserve(n - 1);
  for (int u : s) {
    int lf = ls.top();
    ls.pop();
    es.emplace_back(lf, u);
    deg[lf]--;
    if (--deg[u] == 1) {
      ls.push(u);
    }
  }
  int a = ls.top();
  ls.pop();
  int b = ls.top();
  es.emplace_back(a, b);
  return es;
}

} // namespace noya