utils.h

#pragma once
 
#include <pixie/tree/louds.h>
 
#include <queue>
#include <random>
#include <vector>
 
using pixie::LoudsNode;
 
std::vector<std::vector<size_t>> generate_random_tree(size_t tree_size,
                                                      std::mt19937_64& rng) {
  if (tree_size == 0) {
    return {};
  }
  std::vector<std::vector<size_t>> adj(tree_size);
  adj[0].push_back(0);
  for (size_t i = 1; i < tree_size; i++) {
    size_t parent = rng() % i;
    adj[i].push_back(parent);
    adj[parent].push_back(i);
  }
  return adj;
}
 
std::vector<std::vector<size_t>> bfs_order(
    size_t tree_size,
    const std::vector<std::vector<size_t>>& adj) {
  std::vector<std::vector<size_t>> bfs_adj(tree_size);
  std::queue<std::pair<size_t, size_t>> q;
  bfs_adj[0].push_back(0);
  q.push({0, 0});
  size_t cnt = 1;
  while (!q.empty()) {
    size_t old_v = q.front().first;
    size_t cur_v = q.front().second;
    q.pop();
    for (size_t i = 1; i < adj[old_v].size(); i++) {
      size_t old_u = adj[old_v][i];
      size_t cur_u = cnt++;
      q.push({old_u, cur_u});
      bfs_adj[cur_u].push_back(cur_v);
      bfs_adj[cur_v].push_back(cur_u);
    }
  }
  return bfs_adj;
}
 
std::vector<std::vector<size_t>> dfs_order(
    size_t tree_size,
    const std::vector<std::vector<size_t>>& adj) {
  std::vector<std::vector<size_t>> dfs_adj(tree_size);
  std::vector<std::pair<size_t, size_t>> stack;
  dfs_adj[0].push_back(0);
  stack.push_back({0, 0});
  std::vector<size_t> renumbering(tree_size, 0);
  size_t next_number = 1;
  while (!stack.empty()) {
    auto& [v, i] = stack.back();
    i++;
    if (i == adj[v].size()) {
      stack.pop_back();
      continue;
    }
    size_t u = adj[v][i];
    renumbering[u] = next_number++;
    dfs_adj[renumbering[v]].push_back(renumbering[u]);
    dfs_adj[renumbering[u]].push_back(renumbering[v]);
 
    stack.push_back(std::pair{u, 0});
  }
  return dfs_adj;
}
 
std::vector<uint64_t> adj_to_louds(
    size_t tree_size,
    const std::vector<std::vector<size_t>>& adj) {
  size_t louds_size = tree_size * 2 - 1;
  std::vector<uint64_t> louds((louds_size + 63) / 64, 0);
  size_t pos = 0;
  for (size_t i = 0; i < tree_size; i++) {
    louds[pos >> 6] = louds[pos >> 6] | (1ULL << (pos & 63));
    pos += adj[i].size();
  }
  return louds;
}
 
std::vector<uint64_t> generate_random_data(size_t data_size,
                                           size_t alphabet_size,
                                           std::mt19937_64& rng) {
  std::vector<uint64_t> data(data_size);
  std::uniform_int_distribution<uint64_t> alphabet(0, alphabet_size - 1);
  for (size_t i = 0; i < data_size; i++) {
    data[i] = alphabet(rng);
  }
  return data;
}
 
struct AdjListNode {
  size_t number;
};
 
bool operator==(const AdjListNode& a, const LoudsNode& b) {
  return a.number == b.number;
}
 
bool operator==(const LoudsNode& b, const AdjListNode& a) {
  return a.number == b.number;
}
 
class AdjListTree {
 private:
  std::vector<std::vector<size_t>> adj;
 
 public:
  explicit AdjListTree(const std::vector<std::vector<size_t>>& adjacency_list)
      : adj(adjacency_list) {}

  AdjListNode root() const { return AdjListNode(0); }

  bool is_leaf(const AdjListNode& node) const {
    return adj[node.number].size() <= 1;
  }

  bool is_root(const AdjListNode& node) const { return node.number == 0; }

  size_t degree(const AdjListNode& node) const {
    return adj[node.number].size() - 1;
  }

  AdjListNode child(const AdjListNode& node, size_t i) const {
    return AdjListNode(adj[node.number][i + 1]);
  }

  AdjListNode first_child(const AdjListNode& node) const {
    return AdjListNode(adj[node.number][1]);
  }

  AdjListNode parent(const AdjListNode& node) const {
    return AdjListNode(adj[node.number][0]);
  }

  bool is_last_child(const AdjListNode& node) const {
    size_t p = parent(node).number;
    return adj[p].back() == node.number;
  }

  AdjListNode next_sibling(const AdjListNode& node) const {
    const size_t parent_number = parent(node).number;
    const auto& siblings = adj[parent_number];
    for (size_t i = 1; i + 1 < siblings.size(); ++i) {
      if (siblings[i] == node.number) {
        return AdjListNode(siblings[i + 1]);
      }
    }
    return node;
  }
};