AA Tree key (K) value (V) store single header implementation with subtree sum (S) query in log n (log n pooled allocations)

AATree.hpp

// BSD Zero Clause License

// Copyright (c) 2024 nitko12

// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted.

// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.

// // example usage:

// g++ -std=c++17 -o main main.cpp

// main.cpp

// #include <iostream>
// #include "AATree.hpp"

// int main() {
//   AATree<int, int, int> m;

//   m.insert(1, 5);
//   m.insert(2, 10);
//   m.insert(3, 5);
//   m.insert(4, 10);

//   m.erase(1);

//   std::cout << m.search(1) << std::endl;
//   // 0
//   std::cout << m.search(2) << std::endl;
//   // 1

//   std::cout << m.query_subtree_sum(2, 3) << std::endl;
//   // 15

//   std::cout << m.query_subtree_sum(2, 4) << std::endl;
//   // 25

//   std::cout << m.sum_less_than(3) << std::endl;
//   // 10

//   std::cout << m.to_string() << std::endl;
//   //     4:10[10]
//   // 3:5[25]
//   //   2:10[10]
// }

#pragma once

#include <algorithm>
#include <string>
#include <vector>

template<typename K, typename V, typename S>
struct TreeNode {
  static const size_t NIL = -1;

  K key;
  V value;
  S subtree_sum;

  uint32_t level;
  size_t left, right, parent;

  TreeNode(K key, V value);
};


template<typename K, typename V, typename S>
class AATree {
  public:
  AATree(size_t capacity = 100);

  // core operations
  void insert(K key, V value);
  void erase(K key);
  bool search(K key) const;
  void clear();
  size_t size() const;

  // debugging and validation
  std::string to_string() const;
  bool ensure_subtree_sum() const;
  bool ensure_size() const;

  // range query
  S query_subtree_sum(K L, K R);
  S sum_less_than(K key) const;

  private:
  std::vector<TreeNode<K, V, S>> nodes;
  std::vector<size_t> free_nodes;
  size_t root, size_;

  // helper functions
  size_t insert_(size_t T, K key, V value, size_t parent);
  size_t erase_(size_t T, K key);
  size_t skew_(size_t T);
  size_t split_(size_t T);
  size_t decrease_level_(size_t T);
  size_t find_min_(size_t T) const;
  size_t search_(size_t T, K key) const;
  void clear_(size_t T);
  bool ensure_subtree_sum_(size_t T) const;
  size_t ensure_size_(size_t T) const;
  void to_string_(size_t T, std::string &res, size_t depth) const;

  // utility functions
  void update_subtree_sum(size_t T);
};


// ============================ Implementation ============================

template<typename K, typename V, typename S>
TreeNode<K, V, S>::TreeNode(K key, V value) :
    key(key), value(value), subtree_sum(value), level(1), left(NIL), right(NIL), parent(NIL) {}

template<typename K, typename V, typename S>
AATree<K, V, S>::AATree(size_t capacity) : root(TreeNode<K, V, S>::NIL), size_(0) {
  nodes.reserve(capacity);
  free_nodes.reserve(capacity);
}
template<typename K, typename V, typename S>
void AATree<K, V, S>::update_subtree_sum(size_t T) {
  if (T == TreeNode<K, V, S>::NIL)
    return;
  nodes[T].subtree_sum = nodes[T].value + //
                         (nodes[T].left != TreeNode<K, V, S>::NIL ? nodes[nodes[T].left].subtree_sum : 0) + //
                         (nodes[T].right != TreeNode<K, V, S>::NIL ? nodes[nodes[T].right].subtree_sum : 0);
}
template<typename K, typename V, typename S>
void AATree<K, V, S>::insert(K key, V value) {
  root = insert_(root, key, value, TreeNode<K, V, S>::NIL);
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::insert_(size_t T, K key, V value, size_t parent) {
  if (T == TreeNode<K, V, S>::NIL) {
    size_t idx;
    if (free_nodes.empty()) {
      nodes.emplace_back(key, value);
      idx = nodes.size() - 1;
    } else {
      idx = free_nodes.back();
      free_nodes.pop_back();
      nodes[idx] = TreeNode<K, V, S>(key, value);
    }
    nodes[idx].parent = parent;
    ++size_;
    update_subtree_sum(idx);
    return idx;
  }
  if (key < nodes[T].key)
    nodes[T].left = insert_(nodes[T].left, key, value, T);
  else if (key > nodes[T].key)
    nodes[T].right = insert_(nodes[T].right, key, value, T);

  T = skew_(T);
  T = split_(T);

  update_subtree_sum(T);
  return T;
}
template<typename K, typename V, typename S>
void AATree<K, V, S>::erase(K key) {
  root = erase_(root, key);
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::erase_(size_t T, K key) {
  if (T == TreeNode<K, V, S>::NIL)
    return T;
  if (key < nodes[T].key) {
    nodes[T].left = erase_(nodes[T].left, key);
  } else if (key > nodes[T].key) {
    nodes[T].right = erase_(nodes[T].right, key);
  } else {
    if (nodes[T].left == TreeNode<K, V, S>::NIL || nodes[T].right == TreeNode<K, V, S>::NIL) {
      size_t child = (nodes[T].left != TreeNode<K, V, S>::NIL) ? nodes[T].left : nodes[T].right;
      if (child != TreeNode<K, V, S>::NIL)
        nodes[child].parent = nodes[T].parent;
      --size_;
      free_nodes.push_back(T);
      return child;
    } else {
      size_t successor = find_min_(nodes[T].right);
      nodes[T].key = nodes[successor].key;
      nodes[T].right = erase_(nodes[T].right, nodes[successor].key);
    }
  }

  T = decrease_level_(T);
  T = skew_(T);
  T = split_(T);
  update_subtree_sum(T);
  return T;
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::skew_(size_t T) {
  if (T == TreeNode<K, V, S>::NIL || nodes[T].left == TreeNode<K, V, S>::NIL)
    return T;
  if (nodes[nodes[T].left].level == nodes[T].level) {
    size_t L = nodes[T].left;
    nodes[T].left = nodes[L].right;
    if (nodes[L].right != TreeNode<K, V, S>::NIL)
      nodes[nodes[L].right].parent = T;
    nodes[L].right = T;
    nodes[L].parent = nodes[T].parent;
    nodes[T].parent = L;

    update_subtree_sum(T);
    update_subtree_sum(L);
    return L;
  }
  return T;
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::split_(size_t T) {
  if (T == TreeNode<K, V, S>::NIL || nodes[T].right == TreeNode<K, V, S>::NIL ||
      nodes[nodes[T].right].right == TreeNode<K, V, S>::NIL)
    return T;
  if (nodes[T].level == nodes[nodes[nodes[T].right].right].level) {
    size_t R = nodes[T].right;
    nodes[T].right = nodes[R].left;
    if (nodes[R].left != TreeNode<K, V, S>::NIL)
      nodes[nodes[R].left].parent = T;
    nodes[R].left = T;
    nodes[R].parent = nodes[T].parent;
    nodes[T].parent = R;
    nodes[R].level++;

    update_subtree_sum(T);
    update_subtree_sum(R);
    return R;
  }
  return T;
}
template<typename K, typename V, typename S>
S AATree<K, V, S>::query_subtree_sum(K L, K R) {
  if (R < L)
    return 0;
  return sum_less_than(R + 1) - sum_less_than(L);
}
template<typename K, typename V, typename S>
S AATree<K, V, S>::sum_less_than(K key) const {
  size_t current = root;
  S sum = 0;
  while (current != TreeNode<K, V, S>::NIL) {
    if (nodes[current].key < key) {
      sum += nodes[current].value;
      if (nodes[current].left != TreeNode<K, V, S>::NIL)
        sum += nodes[nodes[current].left].subtree_sum;
      current = nodes[current].right;
    } else {
      current = nodes[current].left;
    }
  }
  return sum;
}
template<typename K, typename V, typename S>
bool AATree<K, V, S>::ensure_size() const {
  return ensure_size_(root) == size_;
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::ensure_size_(size_t T) const {
  if (T == TreeNode<K, V, S>::NIL)
    return 0;
  return 1 + ensure_size_(nodes[T].left) + ensure_size_(nodes[T].right);
}
template<typename K, typename V, typename S>
bool AATree<K, V, S>::ensure_subtree_sum() const {
  return ensure_subtree_sum_(root);
}
template<typename K, typename V, typename S>
bool AATree<K, V, S>::ensure_subtree_sum_(size_t T) const {
  if (T == TreeNode<K, V, S>::NIL)
    return true;
  S sum = nodes[T].value;
  if (nodes[T].left != TreeNode<K, V, S>::NIL)
    sum += nodes[nodes[T].left].subtree_sum;
  if (nodes[T].right != TreeNode<K, V, S>::NIL)
    sum += nodes[nodes[T].right].subtree_sum;
  return sum == nodes[T].subtree_sum && ensure_subtree_sum_(nodes[T].left) && ensure_subtree_sum_(nodes[T].right);
}
template<typename K, typename V, typename S>
inline size_t AATree<K, V, S>::decrease_level_(size_t T) {
  if (T == TreeNode<K, V, S>::NIL)
    return T;
  int32_t should_be = 1 + std::min(nodes[T].left != TreeNode<K, V, S>::NIL ? nodes[nodes[T].left].level : 0,
                                   nodes[T].right != TreeNode<K, V, S>::NIL ? nodes[nodes[T].right].level : 0);
  if (should_be < nodes[T].level) {
    nodes[T].level = should_be;
    if (nodes[T].right != TreeNode<K, V, S>::NIL && should_be < nodes[nodes[T].right].level) {
      nodes[nodes[T].right].level = should_be;
    }
  }
  update_subtree_sum(T);
  return T;
}
template<typename K, typename V, typename S>
bool AATree<K, V, S>::search(K key) const {
  return search_(root, key) != TreeNode<K, V, S>::NIL;
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::search_(size_t T, K key) const {
  while (T != TreeNode<K, V, S>::NIL) {
    if (key < nodes[T].key)
      T = nodes[T].left;
    else if (key > nodes[T].key)
      T = nodes[T].right;
    else
      return T;
  }
  return TreeNode<K, V, S>::NIL;
}
template<typename K, typename V, typename S>
inline size_t AATree<K, V, S>::find_min_(size_t T) const {
  while (nodes[T].left != TreeNode<K, V, S>::NIL)
    T = nodes[T].left;
  return T;
}
template<typename K, typename V, typename S>
std::string AATree<K, V, S>::to_string() const {
  std::string res;
  to_string_(root, res, 0);
  return res;
}
template<typename K, typename V, typename S>
void AATree<K, V, S>::to_string_(size_t T, std::string &res, size_t depth) const {
  if (T == TreeNode<K, V, S>::NIL)
    return;
  to_string_(nodes[T].right, res, depth + 1);
  res += std::string(depth * 2, ' ') + //
         std::to_string(nodes[T].key) + ":" + std::to_string(nodes[T].value) + "[" //
         + std::to_string(nodes[T].subtree_sum) + "]\n";
  to_string_(nodes[T].left, res, depth + 1);
}
template<typename K, typename V, typename S>
void AATree<K, V, S>::clear() {
  clear_(root);
  root = TreeNode<K, V, S>::NIL;
  size_ = 0;
}
template<typename K, typename V, typename S>
void AATree<K, V, S>::clear_(size_t T) {
  if (T == TreeNode<K, V, S>::NIL)
    return;
  clear_(nodes[T].left);
  clear_(nodes[T].right);
  free_nodes.push_back(T);
}
template<typename K, typename V, typename S>
size_t AATree<K, V, S>::size() const {
  return size_;
}