Alex's Anthology of Algorithms Common Code for Contests in Concise C++
Data Structures / Dictionaries and Ordered Sets

2.3.2 AVL Tree

2-Data-Structures/2.3.2_AVL_Tree.cpp

Maintain an ordered map, that is, an ordered collection of key-value pairs such that each possible key appears at most once in the collection. An AVL tree is a binary search tree balanced by height, guaranteeing $O(\log n)$ worst-case running time in insertions and deletions by making sure that the heights of the left and right subtrees at every node differ by at most $1$. Whenever an insertion or deletion breaks this invariant, it is repaired with one or two rotations at each affected node along the search path.

The comparator comp defines the key ordering: comp(a, b) is true when a precedes b. It defaults to std::less<K>; to customize the ordering, instantiate AVLTree<K, V, Compare> and pass the comparator to the constructor.

  • AVLTree<K, V>() constructs an empty map.
  • size() returns the size of the map.
  • empty() returns whether the map is empty.
  • insert(k, v) adds an entry with key k and value v to the map, returning true if a new entry was added or false if the key already exists (in which case the map is unchanged and the old value associated with the key is preserved).
  • erase(k) removes the entry with key k from the map, returning true if the removal was successful or false if the key to be removed was not found.
  • find(k) returns a pointer to a const value associated with key k, or nullptr if the key was not found.
  • entries() returns all key-value entries in comparator order.

The comparator-aware navigation routines min(), max(), lower_bound(k), upper_bound(k), prev(k), and next(k) from the treap in 2.3.1 depend only on the BST property and may be adapted here as needed.

Time Complexity
  • $O(1)$ per call to the constructor, size(), and empty().
  • $O(\log n)$ per call to insert(), erase(), and find(), where $n$ is the number of entries currently in the map.
  • $O(n)$ per call to entries().
Space Complexity
  • $O(n)$ for storage of the map elements.
  • $O(\log n)$ auxiliary stack space for insert(), erase(), entries(), and destruction.
  • $O(n)$ for the vector returned by entries().
  • $O(1)$ auxiliary for all other operations.

Implementation

#include <algorithm>
#include <functional>
#include <utility>
#include <vector>

template<typename K, typename V, typename Compare = std::less<K>>
class AVLTree {
  struct Node {
    K key;
    V value;
    int height;
    Node *left, *right;

    Node(const K &k, const V &v) : key(k), value(v), height(1), left(nullptr), right(nullptr) {}
  } *root;

  int num_nodes;
  Compare comp;

  static int height(Node *n) { return (n != nullptr) ? n->height : 0; }

  static void update_height(Node *n) {
    if (n != nullptr) {
      n->height = 1 + std::max(height(n->left), height(n->right));
    }
  }

  static void rotate_left(Node *&n) {
    Node *tmp = n;
    n = n->right;
    tmp->right = n->left;
    n->left = tmp;
    update_height(tmp);
    update_height(n);
  }

  static void rotate_right(Node *&n) {
    Node *tmp = n;
    n = n->left;
    tmp->left = n->right;
    n->right = tmp;
    update_height(tmp);
    update_height(n);
  }

  static int balance_factor(Node *n) {
    return (n != nullptr) ? (height(n->left) - height(n->right)) : 0;
  }

  static void rebalance(Node *&n) {
    if (n == nullptr) {
      return;
    }
    update_height(n);
    int bf = balance_factor(n);
    if (bf > 1 && balance_factor(n->left) >= 0) {
      rotate_right(n);
    } else if (bf > 1 && balance_factor(n->left) < 0) {
      rotate_left(n->left);
      rotate_right(n);
    } else if (bf < -1 && balance_factor(n->right) <= 0) {
      rotate_left(n);
    } else if (bf < -1 && balance_factor(n->right) > 0) {
      rotate_right(n->right);
      rotate_left(n);
    }
  }

  bool insert(Node *&n, const K &k, const V &v) {
    if (n == nullptr) {
      n = new Node(k, v);
      num_nodes++;
      return true;
    }
    if ((comp(k, n->key) && insert(n->left, k, v)) || (comp(n->key, k) && insert(n->right, k, v))) {
      rebalance(n);
      return true;
    }
    return false;
  }

  bool erase(Node *&n, const K &k) {
    if (n == nullptr) {
      return false;
    }
    if (!(comp(k, n->key) || comp(n->key, k))) {
      if (n->left != nullptr && n->right != nullptr) {
        Node *tmp = n->right;
        while (tmp->left != nullptr) {
          tmp = tmp->left;
        }
        K successor_key = tmp->key;
        n->key = successor_key;
        n->value = tmp->value;
        if (!erase(n->right, successor_key)) {
          return false;
        }
      } else {
        Node *tmp = (n->left != nullptr) ? n->left : n->right;
        delete n;
        n = tmp;
        num_nodes--;
      }
      rebalance(n);
      return true;
    }
    if ((comp(k, n->key) && erase(n->left, k)) || (comp(n->key, k) && erase(n->right, k))) {
      rebalance(n);
      return true;
    }
    return false;
  }

  static void collect_entries(Node *n, std::vector<std::pair<K, V>> &res) {
    if (n != nullptr) {
      collect_entries(n->left, res);
      res.emplace_back(n->key, n->value);
      collect_entries(n->right, res);
    }
  }

  static void clean_up(Node *n) {
    if (n != nullptr) {
      clean_up(n->left);
      clean_up(n->right);
      delete n;
    }
  }

 public:
  explicit AVLTree(Compare comp = Compare{}) : root(nullptr), num_nodes(0), comp(std::move(comp)) {}

  ~AVLTree() { clean_up(root); }
  AVLTree(const AVLTree &) = delete;
  AVLTree &operator=(const AVLTree &) = delete;
  int size() const { return num_nodes; }
  bool empty() const { return root == nullptr; }
  bool insert(const K &k, const V &v) { return insert(root, k, v); }
  bool erase(const K &k) { return erase(root, k); }

  const V *find(const K &k) const {
    Node *n = root;
    while (n != nullptr) {
      if (comp(k, n->key)) {
        n = n->left;
      } else if (comp(n->key, k)) {
        n = n->right;
      } else {
        return &(n->value);
      }
    }
    return nullptr;
  }

  std::vector<std::pair<K, V>> entries() const {
    std::vector<std::pair<K, V>> res;
    res.reserve(num_nodes);
    collect_entries(root, res);
    return res;
  }
};

Example Usage

#include <cassert>
using namespace std;

int main() {
  AVLTree<int, char> t;
  assert(t.empty());
  t.insert(2, 'b');
  t.insert(1, 'a');
  t.insert(3, 'c');
  t.insert(5, 'e');
  assert(t.insert(4, 'd'));
  assert(!t.empty() && t.size() == 5);
  assert(*t.find(4) == 'd');
  assert(!t.insert(4, 'd'));
  assert(t.size() == 5);
  assert(
      (t.entries() == vector<pair<int, char>>{{1, 'a'}, {2, 'b'}, {3, 'c'}, {4, 'd'}, {5, 'e'}})
  );
  assert(t.erase(1));
  assert(!t.erase(1));
  assert(t.find(1) == nullptr);
  assert(t.size() == 4);
  assert((t.entries() == vector<pair<int, char>>{{2, 'b'}, {3, 'c'}, {4, 'd'}, {5, 'e'}}));

  AVLTree<int, int> deep_successor;
  for (int key : {20, 10, 30, 25, 40, 22}) {
    deep_successor.insert(key, key);
  }
  assert(deep_successor.erase(20));
  assert(
      (deep_successor.entries() ==
       vector<pair<int, int>>{{10, 10}, {22, 22}, {25, 25}, {30, 30}, {40, 40}})
  );

  AVLTree<int, char, greater<int>> descending;
  for (int key : {2, 1, 3}) {
    descending.insert(key, '0' + key);
  }
  assert((descending.entries() == vector<pair<int, char>>{{3, '3'}, {2, '2'}, {1, '1'}}));
  assert(*descending.find(2) == '2');
  assert(descending.erase(2) && descending.find(2) == nullptr);
  return 0;
}
Back to top