dfuds.h

#pragma once
 
#include <pixie/rmm/tree.h>
#include <pixie/tree.h>
 
#include <cstddef>
#include <cstdint>
#include <vector>
 
namespace pixie {

template <typename RMMTree>
class DFUDSTree : public TreeBase<DFUDSTree<RMMTree>> {
 private:
  const size_t num_bits_;
  RMMTree rmm_;
 
 public:
  using Node = TreeNode;

  explicit DFUDSTree(const std::vector<std::uint64_t>& dfuds_sequence,
                     size_t tree_size)
      : num_bits_(2 * tree_size - 1), rmm_(dfuds_sequence, 2 * tree_size - 1) {}

  Node root_impl() const { return Node(0, 0); }

  size_t size_impl() const { return (num_bits_ + 1) / 2; }

  bool is_leaf_impl(const Node& node) const {
    return (node.pos + 1 == num_bits_) or rmm_.bit(node.pos) == 0;
  }

  bool is_root_impl(const Node& node) const { return node.number == 0; }

  size_t degree_impl(const Node& node) const {
    return rmm_.select0(node.number + 1) - node.pos;
  }

  Node first_child_impl(const Node& node) const {
    size_t pos = rmm_.select0(node.number + 1);
    size_t num = node.number + 1;
    return Node(num, pos + 1);
  }

  Node child_impl(const Node& node, size_t i) const {
    Node child = first_child_impl(node);
    while (i--) {
      child = next_sibling_impl(child);
    }
    return child;
  }

  Node next_sibling_impl(const Node& node) const {
    size_t end = rmm_.fwdsearch(node.pos, -1);
    size_t pos = end + 1;
    size_t num = rmm_.rank0(pos);
    return Node(num, pos);
  }

  Node parent_impl(const Node& node) const {
    if (node.number == 0) {
      return root_impl();
    }
    size_t open = rmm_.open(
        node.pos -
        1);  // node.pos in 0-based and rmm_.open uses 1-based argument.
             // Thus, we use node.pos meaning the parenthesis before
             // first parenthesis of the current node
    size_t rank = rmm_.rank0(open);
    size_t pos =
        rmm_.select0(rank) +
        1;  // In here we use that rmm_select(0) equals size_t max value so
            // rmm_.select(rank) can still be interpreted as pos-1
    return Node(rank, pos);
  }

  bool is_last_child_impl(const Node& node) const {
    size_t end = rmm_.fwdsearch(node.pos, -1);
    size_t pos = end + 1;
    size_t op = rmm_.open(node.pos - 1);
    size_t op2 = rmm_.open(pos - 1);
    return pos == num_bits_ || op != op2 + 1;
  }
};
 
std::vector<uint64_t> adj_to_dfuds(
    size_t tree_size,
    const std::vector<std::vector<size_t>>& adj) {
  size_t dfuds_size = tree_size * 2 - 1;
  std::vector<uint64_t> dfuds((dfuds_size + 63) / 64, 0);
  std::vector<size_t> stack;
  stack.push_back(0);
  size_t pos = 0;
  while (!stack.empty()) {
    auto v = stack.back();
    stack.pop_back();
    size_t edge_count = adj[v].size();
    for (size_t i = 0; i < edge_count - 1; ++i) {  // edge 0 goes to parent
      dfuds[pos >> 6] = dfuds[pos >> 6] | (1ULL << (pos & 63));
      pos++;
      stack.push_back(adj[v][edge_count - 1 - i]);
    }
    pos++;
  }
  return dfuds;
}
 
}  // namespace pixie