data generation

lib.hpp

#ifndef _DEV_KIWIYOU_GENERATOR_
#define _DEV_KIWIYOU_GENERATOR_

#include <queue>
#include <ranges>
#include <set>
#include <utility>
#include <vector>

namespace kiwiyou {

template<std::ranges::range Range>
std::vector<std::pair<int, int>> tree_from_prufer(Range &&prufer) {
  // Construct edge list of tree from prufer sequence [0, N)^(N - 2).
  std::vector<int> cached_prufer = prufer | std::ranges::to<std::vector>();
  int v_count = std::size(cached_prufer) + 2;
  std::vector<int> degrees(v_count, 0);
  for (auto p : cached_prufer) degrees[p]++;
  std::priority_queue<int> vertices;
  for (auto [u, u_count] : std::views::enumerate(degrees))
    if (u_count == 0) vertices.push(u);
  std::vector<std::pair<int, int>> edges;
  for (auto u : cached_prufer) {
    auto v = vertices.top();
    vertices.pop();
    if (--degrees[u] == 0) vertices.push(u);
    edges.push_back({u, v});
  }
  int u = vertices.top();
  vertices.pop();
  int v = vertices.top();
  edges.push_back({u, v});
  return std::move(edges);
}

template<class IndexByLength>
  requires requires(IndexByLength &&t) { (int)t(0); }
std::vector<int> quicksort_pivot_index_killer(int size, IndexByLength &&pivot) {
  // O(N²) data for pivot fixed by array size
  // returns permutation of [0, N).
  std::vector<int> left, right;
  for (int i = 1; i <= size; i++) {
    if ((int)left.size() < pivot(i)) {
      left.push_back(right.back());
      right.pop_back();
    }
    right.push_back(size - i);
  }
  left.append_range(std::views::reverse(std::move(right)));
  return std::move(left);
}

} // namespace kiwiyou

#endif