Tree.h

//-------------------------------------------------------------------------------------------------------
// Copyright (C) Taylor Woll and panwave contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE.txt file in the project root for
// full license information.
//-------------------------------------------------------------------------------------------------------
 
#ifndef TREE_H
#define TREE_H
 
#include <cassert>
#include <cmath>
#include <vector>
 
namespace panwave {
 
template <class Element, size_t k>
class Tree {
 public:
  explicit Tree(size_t height) : height_(height) {
    assert(height != 0);
 
    this->leaf_count_ = static_cast<size_t>(std::pow(k, height_ - 1));
    size_t const total_nodes = (k * this->leaf_count_ - 1) / (k - 1);
 
    this->nodes_.resize(total_nodes);
    this->mark_.resize(total_nodes);
  }
 
 protected:
  size_t GetLeafCount() const { return this->leaf_count_; }
 
  size_t GetFirstLeaf() const {
    return this->nodes_.size() - this->leaf_count_;
  }
 
  size_t GetLastLeaf() const { return this->nodes_.size() - 1; }
 
  bool IsLeaf(size_t node) const {
    return node >= this->GetFirstLeaf() && node <= this->GetLastLeaf();
  }
 
  size_t GetChild(size_t parent, size_t child_index) const {
    assert(child_index < k);
    assert(parent < this->nodes_.size());
 
    return k * parent + 1 + child_index;
  }
 
  size_t GetParent(size_t child) const {
    assert(child < this->nodes_.size());
    assert(child != 0);
 
    return (child - 1) / k;
  }
 
  Element& GetNodeData(size_t index) {
    assert(index < this->nodes_.size());
 
    return this->nodes_.at(index);
  }
 
  bool IsMarked(size_t node) const {
    assert(node < this->mark_.size());
 
    return this->mark_[node];
  }
 
  void SetMark(size_t node) {
    assert(node < this->mark_.size());
 
    this->mark_[node] = true;
  }
 
  void Unmark() { this->mark_.assign(this->mark_.size(), false); }
 
  size_t GetHeight() const { return this->height_; }
 
 private:
  size_t height_;
  size_t leaf_count_ = 0;
  std::vector<Element> nodes_;
  std::vector<bool> mark_;
};
 
}  // namespace panwave
 
#endif  // TREE_H