Weak AVL Tree

tl;dr: Weak AVL trees are replacements for AVL trees and red-black trees.

The 2014 paper Rank-Balanced Trees (Haeupler, Sen, Tarjan) presents a framework using ranks and rank differences to define binary search trees.

• Each node has a non-negative integer rank r(x). Null nodes have rank -1.
• The rank difference of a node x with parent p(x) is r(p(x)) − r(x).
• A node is i,j if its children have rank differences i and j (unordered), e.g., a 1,2 node has children with rank differences 1 and 2.
• A node is called 1-node if its rank difference is 1.

Several balanced trees fit this framework:

• AVL tree: Ranks are defined as heights. Every node is 1,1 or 1,2 (rank differences of children)
• Red-Black tree: All rank differences are 0 or 1, and no parent of a 0-child is a 0-child. (red: 0-child; black: 1-child; null nodes are black)
• Weak AVL tree (new tree described by this paper): All rank differences are 1 or 2, and every leaf has rank 0.
  • A weak AVL tree without 2,2 nodes is an AVL tree.

AVL trees ⫋ weak AVL trees ⫋ red-black trees

Weak AVL Tree

Weak AVL trees are replacements for AVL trees and red-black trees. A single insertion or deletion operation requires at most two rotations (forming a double rotation when two are needed). In contrast, AVL deletion requires O(log n) rotations, and red-black deletion requires up to three.

Without deletions, a weak AVL tree is exactly an AVL tree. With deletions, its height remains at most that of an AVL tree with the same number of insertions but no deletions.

The rank rules imply:

• Null nodes have rank -1, leaves have rank 0, unary nodes have rank 1.

Insertion

The new node x has a rank of 0, changed from the null node of rank -1. There are three cases.

• If the tree was previously empty, the new node becomes the root.
• If the parent of the new node was previously a unary node (1,2 node), it is now a 1,1 binary node.
• If the parent of the new node was previously a leaf (1,1 node), it is now a 0,1 binary node, leading to a rank violation.

When the tree was previously non-empty, x has a parent node. We call the following subroutine with x indicating the new node to handle the second and third cases.

The following subroutine handles the rank increase of x. We call break if there is no more rank violation, i.e. we are done.

The 2014 paper isn't very clear about the conditions.

// Assume that x's rank has just increased by 1 and rank_diff(x) has been updated.

p = x->parent;
if (rank_diff(x) == 1) {
  // x was previously a 2-child before increasing x->rank.
  // Done.
} else {
  for (;;) {
    // Otherwise, p is a 0,1 node (previously a 1,1 node before increasing x->rank).
    // x being a 0-child is a violation.

    Promote p.
    // Since we have promoted both x and p, it's as if rank_diff(x's sibling) is flipped.
    // p is now a 1,2 node.
    
    x = p;
    p = p->parent;
    // x is a 1,2 node.
    if (!p) break;
    d = p->ch[1] == x;
  
    if (rank_diff(x) == 1) { break; }
    // Otherwise, x is a 0-child, leading to a new rank violation.
    
    auto sib = p->ch[!d];
    if (rank_diff(sib) == 2) { // p is a 0,2 node
      auto y = x->ch[d^1];
      if (y && rank_diff(y) == 1) {
        // y is a 1-child. y must the previous `x` in the last iteration.
        Perform a double rotation involving `p`, `x`, and `y`.
      } else {
        // Otherwise, y is null or a 2-child.
        Perform a single rotation involving `p` and `x`.
        x is now a 1,1 node and there is no more violation.
      }
      break;
    }
  
    // Otherwise, p is a 0,1 node. Goto the next iteration.
  }
}

Insertion never introduces a 2,2 node, so insertion-only sequences produce AVL trees.

Deletion

TODO: Describe deletion

Implementation

In 2020, FreeBSD has changed its sys/tree.h to use weak AVL trees instead of red-black trees. https://reviews.freebsd.org/D25480 The rb_ prefix remains as it can also indicate Rank-Balanced:)

Here is a C++ implementation with the following operations

• insert: insert a node
• remove: remove a node
• rank: count elements less than a key
• select: find the k-th smallest element (0-indexed)
• prev: find the largest element less than a key
• next: find the smallest element greater than a key

The insertion and deletion operations are inspired by the FreeBSD implementation, with the insertion further optimized.

We encode rank differences instead of the absolute rank values. Since valid rank differences can only be 1 or 2, a single bit suffices to encode each. We take 2 bits from the parent pointer to store the two child rank differences.

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <numeric>
#include <random>
#include <vector>
using namespace std;

struct Node {
  Node *ch[2]{};
  uintptr_t par_and_flg{};
  int i{}, sum{}, size{};

  Node *parent() const { return reinterpret_cast<Node*>(par_and_flg & ~3UL); }
  void set_parent(Node *p) { par_and_flg = (par_and_flg & 3) | reinterpret_cast<uintptr_t>(p); }
  uintptr_t flags() const { return par_and_flg & 3; }
  bool rd2(int d) const { return par_and_flg & (1 << d); }
  void flip(int d) { par_and_flg ^= (1 << d); }
  void clr_flags() { par_and_flg &= ~3UL; }

  void mconcat() {
    sum = i;
    size = 1;
    if (ch[0]) sum += ch[0]->sum, size += ch[0]->size;
    if (ch[1]) sum += ch[1]->sum, size += ch[1]->size;
  }

  bool operator<(const Node &o) const { return i < o.i; }
};

struct WAVL {
  Node *root{};

  ~WAVL() {
    auto destroy = [](auto &self, Node *n) -> void {
      if (!n) return;
      self(self, n->ch[0]);
      self(self, n->ch[1]);
      delete n;
    };
    destroy(destroy, root);
  }

  Node *rotate(Node *x, int d) {
    auto pivot = x->ch[d];
    if ((x->ch[d] = pivot->ch[d^1])) x->ch[d]->set_parent(x);
    pivot->set_parent(x->parent());
    if (!x->parent()) root = pivot;
    else x->parent()->ch[x != x->parent()->ch[0]] = pivot;
    pivot->ch[d^1] = x;
    x->set_parent(pivot);
    x->mconcat();
    return pivot;
  }

  void insert(Node *x) {
    Node *p = nullptr;
    int d = 0;
    for (auto tmp = root; tmp; ) {
      p = tmp;
      d = *p < *x;
      tmp = tmp->ch[d];
    }
    x->par_and_flg = reinterpret_cast<uintptr_t>(p);
    x->ch[0] = x->ch[1] = nullptr;
    if (!p) return root = x, x->mconcat();
    p->ch[d] = x;
    auto *x2 = x;
    if (p->rd2(d)) {
      p->flip(d);
    } else {
      assert(p->rd2(d^1) == 0);
      p->flip(d^1);
      int d1 = d;
      for (x = p, p = x->parent(); p; x = p, p = x->parent()) {
        d = (p->ch[1] == x);
        if (p->rd2(d)) {
          p->flip(d);
          break;
        }
        p->flip(d^1);
        if (!p->rd2(d ^ 1)) {
          if ((d^1) == d1) {
            assert(!x->rd2(d1) && (x->ch[d1] == x2 || x->ch[d1]->flags() == 1 || x->ch[d1]->flags() == 2));
            x->flip(d);
            auto y = rotate(x, d^1); // y is previous x
            if (y->rd2(d))
              x->flip(d^1);
            else if (y->rd2(d^1))
              p->flip(d);
            x = y;
          }
          x = rotate(p, d);
          x->clr_flags();
          break;
        }
        d1 = d;
      }
    }
    for (; x2; x2 = x2->parent()) x2->mconcat();
  }

  void remove(Node *x) {
    auto old = x;
    auto p = x->parent();
    auto right = x->ch[1];
    Node *child;
    if (!x->ch[0]) x = child = right;
    else if (!right) x = child = x->ch[0];
    else {
      if (!(child = right->ch[0])) {
        child = right->ch[1];
        p = x = right;
      } else {
        do x = child; while ((child = x->ch[0]));
        child = x->ch[1];
        p = x->parent();
        p->ch[0] = child;
        old->ch[1]->set_parent(x);
        x->ch[1] = old->ch[1];
      }
      old->ch[0]->set_parent(x);
      x->ch[0] = old->ch[0];
      x->par_and_flg = old->par_and_flg;
    }
    if (!old->parent()) root = x;
    else old->parent()->ch[old != old->parent()->ch[0]] = x;
    if (child) child->set_parent(p);

    Node *x2 = p;
    if (p) {
      x = child;
      if (p->ch[0] == x && p->ch[1] == x) {
        p->clr_flags();
        x = p;
        p = x->parent();
      }
      while (p) {
        int d2 = (p->ch[1] == x);
        if (!p->rd2(d2)) {
          p->flip(d2);
          break;
        }
        if (p->rd2(d2 ^ 1)) {
          p->flip(d2 ^ 1);
          x = p;
          p = x->parent();
          continue;
        }
        auto sib = p->ch[d2^1];
        if (sib->flags() == 3) {
          sib->clr_flags();
          x = p;
          p = x->parent();
          continue;
        }
        sib->flip(d2^1);
        if (sib->rd2(d2))
          p->flip(d2);
        else if (!sib->rd2(d2^1)) {
          p->flip(d2);
          x = rotate(sib, d2);
          if (x->rd2(d2^1)) sib->flip(d2);
          if (x->rd2(d2)) p->flip(d2^1);
          x->par_and_flg |= 3;
        }
        rotate(p, d2^1);
        break;
      }
    }
    for (; x2; x2 = x2->parent()) x2->mconcat();
  }

  Node *find(int key) const {
    auto tmp = root;
    while (tmp) {
      if (key < tmp->i) tmp = tmp->ch[0];
      else if (key > tmp->i) tmp = tmp->ch[1];
      else return tmp;
    }
    return nullptr;
  }

  Node *min() const {
    Node *p = nullptr;
    for (auto n = root; n; n = n->ch[0]) p = n;
    return p;
  }

  int rank(int key) const {
    int r = 0;
    for (auto n = root; n; ) {
      if (key <= n->i) n = n->ch[0];
      else {
        r += 1 + (n->ch[0] ? n->ch[0]->size : 0);
        n = n->ch[1];
      }
    }
    return r;
  }

  int select(int k) const {
    auto x = root;
    while (x) {
      int lsz = x->ch[0] ? x->ch[0]->size : 0;
      if (k < lsz) x = x->ch[0];
      else if (k == lsz) return x->i;
      else k -= lsz + 1, x = x->ch[1];
    }
    return -1;
  }

  int prev(int key) const {
    int res = -1;
    for (auto x = root; x; )
      if (key <= x->i) x = x->ch[0];
      else { res = x->i; x = x->ch[1]; }
    return res;
  }

  int next(int key) const {
    int res = -1;
    for (auto x = root; x; )
      if (key >= x->i) x = x->ch[1];
      else { res = x->i; x = x->ch[0]; }
    return res;
  }

  static Node *next(Node *x) {
    if (x->ch[1]) {
      x = x->ch[1];
      while (x->ch[0]) x = x->ch[0];
    } else {
      while (x->parent() && x == x->parent()->ch[1]) x = x->parent();
      x = x->parent();
    }
    return x;
  }
};

void print_tree(Node *n, int d = 0) {
  if (!n) { printf("%*snil\n", 2*d, ""); return; }
  print_tree(n->ch[0], d + 1);
  printf("%*s%d (%d,%d)\n", 2*d, "", n->i, n->rd2(0) ? 2 : 1, n->rd2(1) ? 2 : 1);
  print_tree(n->ch[1], d + 1);
}

int compute_rank(Node *n, bool debug = false) {
  if (!n) return -1;
  int lr = compute_rank(n->ch[0], debug), rr = compute_rank(n->ch[1], debug);
  if (lr < -1 || rr < -1) return -2;
  int rank_l = lr + (n->rd2(0) ? 2 : 1);
  int rank_r = rr + (n->rd2(1) ? 2 : 1);
  if (rank_l != rank_r) {
    if (debug) printf("node %d: rank mismatch left=%d right=%d\n", n->i, rank_l, rank_r);
    return -2;
  }
  if (!n->ch[0] && !n->ch[1] && n->flags() != 0) {
    if (debug) printf("node %d: leaf must be 1,1 but flags=%lu\n", n->i, n->flags());
    return -2;
  }
  int expected_sum = n->i + (n->ch[0] ? n->ch[0]->sum : 0) + (n->ch[1] ? n->ch[1]->sum : 0);
  if (n->sum != expected_sum) {
    if (debug) printf("node %d: sum mismatch got=%d expected=%d\n", n->i, n->sum, expected_sum);
    return -2;
  }
  int expected_size = 1 + (n->ch[0] ? n->ch[0]->size : 0) + (n->ch[1] ? n->ch[1]->size : 0);
  if (n->size != expected_size) {
    if (debug) printf("node %d: size mismatch got=%d expected=%d\n", n->i, n->size, expected_size);
    return -2;
  }
  return rank_l;
}

bool verify_tree(const WAVL &tree, bool verbose = false) {
  int rank = compute_rank(tree.root);
  if (rank < -1) {
    printf("INVALID TREE\n");
    compute_rank(tree.root, true);
    return false;
  }
  if (verbose) printf("Tree verified, rank = %d\n", rank);
  return true;
}

int main() {
  srand(42);
  WAVL tree;
  int i = 0;
  std::vector<int> a(20);
  std::iota(a.begin(), a.end(), 1);
  std::shuffle(a.begin(), a.end(), std::default_random_engine(42));
  for (int val : a) {
    auto n = new Node;
    n->i = val;
    tree.insert(n);
    if (i++ < 6) {
      printf("-- %d After insertion of %d\n", i, val);
      print_tree(tree.root);
    }
  }
  printf("\nSum\tof values = %d\n", tree.root->sum);
  verify_tree(tree, true);

  for (int val : {5, 10, 15}) {
    if (auto found = tree.find(val)) {
      tree.remove(found);
      delete found;
    }
  }
  printf("After removing 5, 10, 15:\n");
  printf("\nSum\tof values = %d\n", tree.root->sum);
  verify_tree(tree, true);

  std::vector<Node*> ref;
  for (auto n = tree.min(); n; n = WAVL::next(n)) ref.push_back(n);

  for (int i = 0; i < 100000; i++) {
    if (ref.size() < 5 || (ref.size() < 1000 && rand() % 2 == 0)) {
      auto n = new Node;
      n->i = rand() % 100000;
      tree.insert(n);
      ref.push_back(n);
    } else {
      int idx = rand() % ref.size();
      tree.remove(ref[idx]);
      delete ref[idx];
      ref[idx] = ref.back();
      ref.pop_back();
    }
    if (i%100 == 0 && !verify_tree(tree)) {
      printf("FAILED at iteration %d\n", i);
      return 1;
    }
  }

  while (!ref.empty()) {
    tree.remove(ref.back());
    delete ref.back();
    ref.pop_back();
    if (tree.root && !verify_tree(tree)) {
      printf("FAILED during final cleanup\n");
      return 1;
    }
  }
  printf("Stress test passed\n");

  // Test rank, select, prev, next
  printf("\nTesting rank/select/prev/next...\n");
  std::vector<int> vals = {10, 20, 30, 40, 50};
  for (int v : vals) {
    auto n = new Node;
    n->i = v;
    tree.insert(n);
  }

  // rank tests (number of elements < key)
  assert(tree.rank(5) == 0);
  assert(tree.rank(10) == 0);
  assert(tree.rank(15) == 1);
  assert(tree.rank(20) == 1);
  assert(tree.rank(25) == 2);
  assert(tree.rank(50) == 4);
  assert(tree.rank(55) == 5);

  // select tests (0-indexed)
  assert(tree.select(0) == 10);
  assert(tree.select(1) == 20);
  assert(tree.select(2) == 30);
  assert(tree.select(3) == 40);
  assert(tree.select(4) == 50);
  assert(tree.select(5) == -1);

  // prev tests (largest < key)
  assert(tree.prev(10) == -1);
  assert(tree.prev(11) == 10);
  assert(tree.prev(20) == 10);
  assert(tree.prev(21) == 20);
  assert(tree.prev(50) == 40);
  assert(tree.prev(55) == 50);

  // next tests (smallest > key)
  assert(tree.next(5) == 10);
  assert(tree.next(10) == 20);
  assert(tree.next(15) == 20);
  assert(tree.next(40) == 50);
  assert(tree.next(50) == -1);
  assert(tree.next(55) == -1);

  printf("rank/select/prev/next tests passed\n");
}

Node structure:

• ch[2]: left and right child pointers.
• par_and_flg: packs the parent pointer with 2 flag bits in the low bits. Bit 0 indicates whether the left child has rank difference 2; bit 1 indicates whether the right child has rank difference 2. A cleared bit means rank difference 1.
• i: the key value.
• sum, size: augmented data maintained by mconcat for order statistics operations.

Helper methods:

• rd2(d): returns true if child d has rank difference 2.
• flip(d): toggles the rank difference of child d between 1 and 2.
• clr_flags(): sets both children to rank difference 1 (used after rotations to reset a node to 1,1).

Invariants:

• Leaves always have flags() == 0, meaning both null children are 1-children (null nodes have rank -1, leaf has rank 0).
• After each insertion or deletion, mconcat is called along the path to the root to update augmented data.

Rotations:

The rotate(x, d) function rotates node x in direction d. It lifts x->ch[d] to replace x, and updates the augmented data for x. The caller is responsible for updating rank differences.

The paper suggests that each node can use one bit to encode its own rank difference. However, this representation is not ideal. A null node can be either a 1-child (parent is binary) or a 2-child (parent is unary). Retrieving the child rank difference requires probing the sibling node:

int rank_diff(Node *p, int d) { return p->ch[d] ? p->ch[d]->par_and_flg & 1 : p->ch[!d] ? 2 : 1; }

Misc

Visualization: https://tjkendev.github.io/bst-visualization/avl-tree/bu-weak.html