dfuds_tree.h

#pragma once
 
#include <pixie/rmm_tree.h>
 
#include <cstdint>
 
#include "utils.h"
 
namespace pixie {

class DFUDSTree {
 private:
  const size_t num_bits_;
  RmMTree rmm_;
 
 public:
  struct Node {
    size_t number;
 
    size_t pos;

    Node(size_t node_number, size_t dfuds_pos)
        : number(node_number), pos(dfuds_pos) {}
  };

  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) {}

  static Node root() { return Node(0, 0); }

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

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

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

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

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

  Node child(const Node& node, size_t i) const {
    size_t pos = rmm_.close(rmm_.select0(node.number + 1) - i) + 1;
    size_t num = rmm_.rank0(pos);
    return Node(num, pos);
  }

  Node next_sibling(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(const Node& node) const {
    if (node.number == 0) {
      return root();
    }
    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(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;
}
 
bool operator==(const AdjListNode& a, const DFUDSTree::Node& b) {
  return a.number == b.number;
}
 
bool operator==(const DFUDSTree::Node& b, const AdjListNode& a) {
  return a.number == b.number;
}
 
}  // namespace pixie