avl_order_statistic_set.hpp

#ifndef AVL_ORDER_STATISTIC_SET_HPP
#define AVL_ORDER_STATISTIC_SET_HPP

#include <algorithm>
#include <cstddef>
#include <functional>
#include <initializer_list>
#include <utility>

namespace avl_order_statistic_set {
/**
 * @brief Node structure for AVL Order Statistic Tree
 *
 * @tparam T Type of elements stored in the node
 */
template <typename T> struct Node {
  T key;                  ///< The key/value stored in the node
  Node *left = nullptr;   ///< Pointer to left child
  Node *right = nullptr;  ///< Pointer to right child
  Node *parent = nullptr; ///< Pointer to parent node
  int height = 1;         ///< Height of the subtree rooted at this node
  std::size_t size = 1;   ///< Size of the subtree (for order statistics)

  /**
   * @brief Construct a new Node object (lvalue key version)
   *
   * @param value The key value to store
   * @param p Pointer to parent node (default nullptr)
   */
  explicit Node(const T &value, Node *p = nullptr) : key(value), parent(p) {}

  /**
   * @brief Construct a new Node object (rvalue key version)
   *
   * @param value The key value to store (will be moved)
   * @param p Pointer to parent node (default nullptr)
   */
  explicit Node(T &&value, Node *p = nullptr)
      : key(std::move(value)), parent(p) {}

  // Disallow copying
  Node(const Node &) = delete;
  Node &operator=(const Node &) = delete;
};

// ==================== UTILITY FUNCTIONS ====================

/**
 * @brief Get the height of a node (nullptr-safe)
 *
 * @tparam T Type of node elements
 * @param cur Node to check
 * @return int Height of the node (0 if nullptr)
 */
template <typename T> int get_height(const Node<T> *cur) noexcept {
  return cur ? cur->height : 0;
}

/**
 * @brief Get the subtree size of a node (nullptr-safe)
 *
 * @tparam T Type of node elements
 * @param cur Node to check
 * @return std::size_t Size of subtree (0 if nullptr)
 */
template <typename T>
std::size_t get_subtree_size(const Node<T> *cur) noexcept {
  return cur ? cur->size : 0;
}

/**
 * @brief Calculate balance factor of a node
 *
 * @tparam T Type of node elements
 * @param cur Node to check
 * @return int Balance factor (left height - right height)
 */
template <typename T>
int calculate_balance_factor(const Node<T> *cur) noexcept {
  return get_height(cur->left) - get_height(cur->right);
}

/**
 * @brief Update node metadata (height and size)
 *
 * @tparam T Type of node elements
 * @param cur Node to update
 */
template <typename T> void update_node_metadata(Node<T> *cur) noexcept {
  if (cur) {
    cur->height = 1 + std::max(get_height(cur->left), get_height(cur->right));
    cur->size = 1 + get_subtree_size(cur->left) + get_subtree_size(cur->right);
  }
}

// ==================== TREE ROTATIONS ====================

/**
 * @brief Perform left rotation around a node
 *
 * @tparam T Type of node elements
 * @param x The pivot node
 * @return Node<T>* The new root of this subtree
 */
template <typename T> Node<T> *rotate_left(Node<T> *x) noexcept {
  Node<T> *y = x->right;
  x->right = y->left;
  if (y->left)
    y->left->parent = x;

  y->parent = x->parent;
  y->left = x;
  x->parent = y;

  update_node_metadata(x);
  update_node_metadata(y);
  return y;
}

/**
 * @brief Perform right rotation around a node
 *
 * @tparam T Type of node elements
 * @param y The pivot node
 * @return Node<T>* The new root of this subtree
 */
template <typename T> Node<T> *rotate_right(Node<T> *y) noexcept {
  Node<T> *x = y->left;
  y->left = x->right;
  if (x->right)
    x->right->parent = y;

  x->parent = y->parent;
  x->right = y;
  y->parent = x;

  update_node_metadata(y);
  update_node_metadata(x);
  return x;
}

/**
 * @brief Rebalance the tree at the given node if needed
 *
 * @tparam T Type of node elements
 * @param cur Node to rebalance
 * @return Node<T>* New root of the subtree after rebalancing
 */
template <typename T> Node<T> *rebalance(Node<T> *cur) noexcept {
  update_node_metadata(cur);
  int b = calculate_balance_factor(cur);

  // Left heavy
  if (b > 1) {
    if (calculate_balance_factor(cur->left) < 0)
      cur->left = rotate_left(cur->left);
    cur = rotate_right(cur);
  }
  // Right heavy
  else if (b < -1) {
    if (calculate_balance_factor(cur->right) > 0)
      cur->right = rotate_right(cur->right);
    cur = rotate_left(cur);
  }

  update_node_metadata(cur);
  return cur;
}

// ==================== TREE OPERATIONS ====================

/**
 * @brief Find the node with minimum key in a subtree
 *
 * @tparam T Type of node elements
 * @param cur Root of the subtree
 * @return Node<T>* Pointer to minimum node
 */
template <typename T> Node<T> *find_min_node(Node<T> *cur) noexcept {
  while (cur && cur->left)
    cur = cur->left;
  return cur;
}

/**
 * @brief Find the node with maximum key in a subtree
 *
 * @tparam T Type of node elements
 * @param cur Root of the subtree
 * @return Node<T>* Pointer to maximum node
 */
template <typename T> Node<T> *find_max_node(Node<T> *cur) noexcept {
  while (cur && cur->right)
    cur = cur->right;
  return cur;
}

/**
 * @brief Recursively destroy a subtree
 *
 * @tparam T Type of node elements
 * @param cur Root of the subtree to destroy
 */
template <typename T> void destroy(Node<T> *cur) noexcept {
  if (cur) {
    destroy(cur->left);
    destroy(cur->right);
    delete cur;
  }
}

/**
 * @brief Deep copy a subtree
 *
 * @tparam T Type of node elements
 * @param src Root of source subtree
 * @param parent Parent node for the new subtree
 * @return Node<T>* Root of the new subtree
 */
template <typename T> Node<T> *deep_copy(Node<T> *src, Node<T> *parent) {
  if (!src)
    return nullptr;

  Node<T> *n = new Node<T>(src->key, parent);
  n->height = src->height;
  n->size = src->size;
  n->left = deep_copy(src->left, n);
  n->right = deep_copy(src->right, n);

  return n;
}

// ==================== ITERATOR TRAITS ====================

template <typename T> struct AVLOrderStatisticSetIteratorTraits {
  using iterator_category = std::bidirectional_iterator_tag;
  using value_type = T;
  using difference_type = std::ptrdiff_t;
  using pointer = const T *;
  using reference = const T &;
};

// ==================== MAIN CLASS ====================

/**
 * @brief AVL Tree based Order Statistic Set
 *
 * @tparam T Type of elements to store
 * @tparam Compare Comparison function object type (default std::less<T>)
 *
 * This implementation provides:
 * - O(log n) insertion, deletion and search
 * - O(log n) order statistics operations
 * - Full STL-style iterator support
 * - Exception safety basic guarantee
 */
template <typename T, typename Compare = std::less<T>>
class AVLOrderStatisticSet {
  Node<T> *root = nullptr; ///< Root of the AVL tree
  std::size_t n = 0;       ///< Number of elements in the set
  Compare comp;            ///< Comparison function object

  // ==================== PRIVATE HELPERS ====================

  /**
   * @brief Find a node with given key
   *
   * @param cur Current node in search
   * @param key Key to search for
   * @return Node<T>* Pointer to found node or nullptr
   */
  Node<T> *find_helper(Node<T> *cur, const T &key) const noexcept {
    while (cur) {
      if (comp(key, cur->key))
        cur = cur->left;
      else if (comp(cur->key, key))
        cur = cur->right;
      else
        return cur;
    }
    return nullptr;
  }

  /**
   * @brief Insert helper (lvalue key version) that returns both new root and
   * inserted node
   *
   * @param cur Current node in recursion
   * @param key Key to insert
   * @param parent Parent node
   * @param inserted [out] Set to true if insertion occurred
   * @return std::pair<Node<T>*, Node<T>*> Pair of {new root, inserted node}
   */
  std::pair<Node<T> *, Node<T> *>
  insert_helper(Node<T> *cur, const T &key, Node<T> *parent, bool &inserted) {
    if (!cur) {
      Node<T> *nn = new Node<T>(key, parent);

      inserted = true;
      ++n;

      return {nn, nn};
    }

    Node<T> *inserted_node = nullptr;
    if (comp(key, cur->key)) {
      auto result = insert_helper(cur->left, key, cur, inserted);
      cur->left = result.first;
      inserted_node = result.second;
    } else if (comp(cur->key, key)) {
      auto result = insert_helper(cur->right, key, cur, inserted);
      cur->right = result.first;
      inserted_node = result.second;
    } else {
      // Key already exists, return current node as inserted node
      inserted_node = cur;
    }

    return {rebalance(cur), inserted_node};
  }

  /**
   * @brief Insert helper (rvalue key version) that returns both new root and
   * inserted node
   *
   * @param cur Current node in recursion
   * @param key Key to insert (will be moved)
   * @param parent Parent node
   * @param inserted [out] Set to true if insertion occurred
   * @return std::pair<Node<T>*, Node<T>*> Pair of {new root, inserted node}
   */
  std::pair<Node<T> *, Node<T> *>
  insert_helper(Node<T> *cur, T &&key, Node<T> *parent, bool &inserted) {
    if (!cur) {
      Node<T> *nn = new Node<T>(std::forward<T>(key), parent);

      inserted = true;
      ++n;

      return {nn, nn};
    }

    Node<T> *inserted_node = nullptr;
    if (comp(key, cur->key)) {
      auto result =
          insert_helper(cur->left, std::forward<T>(key), cur, inserted);
      cur->left = result.first;
      inserted_node = result.second;
    } else if (comp(cur->key, key)) {
      auto result =
          insert_helper(cur->right, std::forward<T>(key), cur, inserted);
      cur->right = result.first;
      inserted_node = result.second;
    } else {
      // Key already exists, return current node as inserted node
      inserted_node = cur;
    }

    return {rebalance(cur), inserted_node};
  }

  /**
   * @brief Erase helper
   *
   * @param cur Current node in recursion
   * @param key Key to erase
   * @param erased [out] Set to true if erasure occurred
   * @return Node<T>* New root of subtree
   */
  Node<T> *erase_helper(Node<T> *cur, const T &key, bool &erased) {
    if (!cur)
      return nullptr;

    if (comp(key, cur->key)) {
      cur->left = erase_helper(cur->left, key, erased);
    } else if (comp(cur->key, key)) {
      cur->right = erase_helper(cur->right, key, erased);
    } else {
      // Actual deletion---set 'erased' and decrement n here ONLY
      erased = true;
      --n;

      // Node with only one child or no child
      if (!cur->left) {
        Node<T> *right = cur->right;
        if (right)
          right->parent = cur->parent;
        delete cur;
        return right;
      } else if (!cur->right) {
        Node<T> *left = cur->left;
        if (left)
          left->parent = cur->parent;
        delete cur;
        return left;
      }
      // Node with two children
      else {
        Node<T> *y = find_min_node(cur->right);
        cur->key = std::move(y->key);
        bool dummy = false; // Don't double-count
        cur->right = erase_helper(cur->right, cur->key, dummy);
      }
    }

    return rebalance(cur);
  }

  /**
   * @brief Erase a node given a direct pointer to the node (used for erasing by
   * iterator).
   *
   * @tparam T Type of node elements
   * @param cur The root of the subtree in which to find and erase the node
   * @param to_delete Pointer to the node to erase
   * @param erased [out] Set to true if erasure occurred
   * @return Node<T>* New root of the subtree after erasure
   *
   * Removes the specified node from the tree and rebalances as necessary.
   */
  Node<T> *erase_node_helper(Node<T> *cur, Node<T> *to_delete, bool &erased) {
    if (!cur)
      return nullptr;

    if (cur == to_delete) {
      // Actual deletion---set 'erased' and decrement n here ONLY
      erased = true;
      --n;

      // Node with only one child or no child
      if (!cur->left) {
        Node<T> *right = cur->right;
        if (right)
          right->parent = cur->parent;
        delete cur;
        return right;
      } else if (!cur->right) {
        Node<T> *left = cur->left;
        if (left)
          left->parent = cur->parent;
        delete cur;
        return left;
      }
      // Node with two children
      else {
        Node<T> *y = find_min_node(cur->right);
        cur->key = std::move(y->key);
        bool dummy = false; // Don't double-count
        cur->right = erase_helper(cur->right, cur->key, dummy);
      }
    } else if (comp(to_delete->key, cur->key)) {
      cur->left = erase_node_helper(cur->left, to_delete, erased);
    } else {

      cur->right = erase_node_helper(cur->right, to_delete, erased);
    }

    return rebalance(cur);
  }

public:
  // ==================== CONSTRUCTORS/DESTRUCTOR ====================

  /// Default constructor
  AVLOrderStatisticSet() noexcept = default;

  template <typename InputIt>
  AVLOrderStatisticSet(InputIt first, InputIt last) {
    for (; first != last; ++first)
      insert(*first);
  }

  /**
   * @brief Construct from initializer list
   *
   * @param il Initializer list of elements
   */
  AVLOrderStatisticSet(std::initializer_list<T> il) {
    for (const auto &v : il)
      insert(v);
  }

  /// Move constructor
  AVLOrderStatisticSet(AVLOrderStatisticSet<T, Compare> &&other) noexcept
      : root(other.root), n(other.n), comp(std::move(other.comp)) {
    other.root = nullptr;
    other.n = 0;
  }

  /// Move assignment
  AVLOrderStatisticSet<T, Compare> &
  operator=(AVLOrderStatisticSet<T, Compare> &&other) noexcept {
    if (this != &other) {
      clear();
      root = other.root;
      n = other.n;
      comp = std::move(other.comp);
      other.root = nullptr;
      other.n = 0;
    }
    return *this;
  }

  /// Copy constructor
  AVLOrderStatisticSet(const AVLOrderStatisticSet<T, Compare> &other)
      : n(other.n) {
    root = deep_copy(other.root, nullptr);
  }

  /// Copy assignment
  AVLOrderStatisticSet<T, Compare> &
  operator=(const AVLOrderStatisticSet<T, Compare> &other) {
    if (this != &other) {
      clear();
      n = other.n;
      root = deep_copy(other.root, nullptr);
    }
    return *this;
  }

  /// Destructor
  ~AVLOrderStatisticSet() { clear(); }

  // ==================== ITERATOR IMPLEMENTATION ====================

  /**
   * @brief Bidirectional iterator for AVLOrderStatisticSet
   */
  class iterator {
    Node<T> *node = nullptr;                                ///< Current node
    const AVLOrderStatisticSet<T, Compare> *tree = nullptr; ///< Associated tree

    iterator(Node<T> *n, const AVLOrderStatisticSet<T, Compare> *t)
        : node(n), tree(t) {}

  public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = T;
    using difference_type = std::ptrdiff_t;
    using pointer = const T *;
    using reference = const T &;

    iterator() = default;
    iterator(const iterator &) = default;
    iterator &operator=(const iterator &) = default;

    reference operator*() const { return node->key; }
    pointer operator->() const { return &node->key; }

    bool operator==(const iterator &other) const { return node == other.node; }
    bool operator!=(const iterator &other) const { return node != other.node; }

    /// Pre-increment (move to next larger key)
    iterator &operator++() {
      if (!node)
        return *this;

      if (node->right) {
        node = find_min_node(node->right);
      } else {
        Node<T> *p = node->parent;
        while (p && node == p->right) {
          node = p;
          p = p->parent;
        }
        node = p;
      }
      return *this;
    }

    /// Post-increment
    iterator operator++(int) {
      iterator tmp(*this);
      ++(*this);
      return tmp;
    }

    /// Pre-decrement (move to next smaller key)
    iterator &operator--() {
      if (!node) {
        node = find_max_node(tree->root);
      } else if (node->left) {
        node = find_max_node(node->left);
      } else {
        Node<T> *p = node->parent;
        while (p && node == p->left) {
          node = p;
          p = p->parent;
        }
        node = p;
      }
      return *this;
    }

    /// Post-decrement
    iterator operator--(int) {
      iterator tmp(*this);
      --(*this);
      return tmp;
    }

    friend class AVLOrderStatisticSet;
  };

  // Iterator type aliases
  using reverse_iterator = std::reverse_iterator<iterator>;
  using const_iterator = iterator;
  using const_reverse_iterator = std::reverse_iterator<const_iterator>;

  // ==================== ITERATOR ACCESS ====================

  iterator begin() const { return iterator(find_min_node(root), this); }
  iterator end() const { return iterator(nullptr, this); }
  reverse_iterator rbegin() const { return reverse_iterator(end()); }
  reverse_iterator rend() const { return reverse_iterator(begin()); }
  const_iterator cbegin() const { return begin(); }
  const_iterator cend() const { return end(); }
  const_reverse_iterator crbegin() const {
    return const_reverse_iterator(cend());
  }
  const_reverse_iterator crend() const {
    return const_reverse_iterator(cbegin());
  }
  friend iterator begin(const AVLOrderStatisticSet<T, Compare> &s) {
    return s.begin();
  }
  friend iterator end(const AVLOrderStatisticSet<T, Compare> &s) {
    return s.end();
  }

  // ==================== EQUALITY ====================

  /**
   * @brief Checks if two AVLOrderStatisticSets are equal
   *
   * Two sets are considered equal if they contain the same elements.
   * This is determined by:
   * 1. Comparing sizes (must be equal)
   * 2. Comparing elements in order (since sets are ordered)
   *
   * @param other The set to compare with
   * @return true If both sets contain exactly the same elements
   * @return false If sets differ in size or element content
   */
  bool operator==(const AVLOrderStatisticSet<T, Compare> &other) const {
    return size() == other.size() && std::equal(begin(), end(), other.begin());
  }

  // ==================== CAPACITY ====================

  std::size_t size() const noexcept { return n; }
  bool empty() const noexcept { return n == 0; }
  explicit operator bool() const noexcept { return !empty(); }

  // ==================== MODIFIERS ====================

  /// Clear all elements from the set
  void clear() noexcept {
    destroy(root);
    root = nullptr;
    n = 0;
  }

  /**
   * @brief Insert a key (lvalue version)
   *
   * @param key Key to insert
   * @return std::pair<iterator, bool> Iterator to inserted element and success
   * flag
   */
  std::pair<iterator, bool> insert(const T &key) {
    bool inserted = false;
    auto result = insert_helper(root, key, nullptr, inserted);
    root = result.first;
    return {iterator(result.second, this), inserted};
  }

  /**
   * @brief Insert a key (rvalue version)
   *
   * @param key Key to insert (will be moved)
   * @return std::pair<iterator, bool> Iterator to inserted element and success
   * flag
   */
  std::pair<iterator, bool> insert(T &&key) {
    bool inserted = false;
    auto result = insert_helper(root, std::move(key), nullptr, inserted);
    root = result.first;
    return {iterator(result.second, this), inserted};
  }

  /**
   * @brief Constructs element in-place if the key doesn't exist
   *
   * @tparam Args Types of arguments to construct the element
   * @param args Arguments to forward to the constructor of the element
   * @return std::pair<iterator, bool>
   *         pair::first: Iterator to the inserted element or the element that
   * prevented insertion pair::second: true if insertion occurred, false
   * otherwise
   */
  template <typename... Args>
  std::pair<iterator, bool> emplace(Args &&...args) {
    T key(std::forward<Args>(args)...);
    return insert(std::move(key));
  }

  /**
   * @brief Erases element with the specified key if it exists
   *
   * @param key Key of the element to erase
   * @return std::size_t Number of elements erased (0 or 1)
   */
  std::size_t erase(const T &key) {
    bool erased = false;
    root = erase_helper(root, key, erased);
    return erased ? 1 : 0;
  }

  /**
   * @brief Erase by iterator: returns iterator to next element
   *
   * @param pos Iterator of the element to erase
   * @return iterator Iterator to next element
   */
  iterator erase(iterator pos) {
    if (pos == end())
      return end();

    Node<T> *del_node = pos.node;

    iterator next = pos;
    ++next;

    bool erased = false;
    root = erase_node_helper(root, del_node, erased);

    return next;
  }

  /**
   * @brief Swaps the contents of this set with another.
   *
   * @param other The other AVLOrderStatisticSet to swap with
   *
   * Swaps the roots, element counts, and comparison functors of the two sets.
   */
  void swap(AVLOrderStatisticSet<T, Compare> &other) noexcept {
    std::swap(root, other.root);
    std::swap(n, other.n);
    std::swap(comp, other.comp);
  }

  // ==================== LOOKUP ====================

  /**
   * @brief Finds an element with key equivalent to key
   *
   * @param key Key value to search for
   * @return iterator Iterator to the found element, or end() if not found
   */
  iterator find(const T &key) const {
    return iterator(find_helper(root, key), this);
  }

  /**
   * @brief Returns the number of elements with key equivalent to key
   *
   * @param key Key value to count
   * @return std::size_t 1 if the container contains the key, 0 otherwise
   *
   * Since this is a set, the count will always be 0 or 1
   */
  std::size_t count(const T &key) const {
    return find_helper(root, key) ? 1 : 0;
  }

  /**
   * @brief Finds an iterator to the first element not less than (i.e. greater
   * or equal to) the specified key.
   *
   * @param key The key to compare against
   * @return iterator Iterator to the first element that is not less than key,
   *                 or end() if every element in the set is less than key.
   *
   * Complexity: O(log n)
   */
  iterator lower_bound(const T &key) const {
    Node<T> *cur = root, *res = nullptr;

    while (cur) {
      if (!comp(cur->key, key)) {
        res = cur;
        cur = cur->left;
      } else
        cur = cur->right;
    }

    return iterator(res, this);
  }

  /**
   * @brief Finds an iterator to the first element greater than the specified
   * key.
   *
   * @param key The key to compare against
   * @return iterator Iterator to the first element that is greater than key,
   *                 or end() if no such element exists.
   *
   * Complexity: O(log n)
   */
  iterator upper_bound(const T &key) const {
    Node<T> *cur = root, *res = nullptr;

    while (cur) {
      if (comp(key, cur->key)) {
        res = cur;
        cur = cur->left;
      } else
        cur = cur->right;
    }

    return iterator(res, this);
  }

  // ==================== ORDER STATISTICS ====================

  /**
   * @brief Finds the element at the specified position in the sorted sequence
   *
   * @param k 0-based index of the element to find
   * @return iterator Iterator to the element at position k, or end() if k >=
   * size()
   *
   * Complexity: O(log n)
   */
  iterator find_by_order(std::size_t k) const {
    Node<T> *cur = root;
    while (cur) {
      std::size_t left_sz = get_subtree_size(cur->left);
      if (k < left_sz)
        cur = cur->left;
      else if (k > left_sz) {
        k -= left_sz + 1;
        cur = cur->right;
      } else {
        return iterator(cur, this);
      }
    }
    return end(); // not found, k >= size()
  }

  /**
   * @brief Returns the number of elements with keys less than key
   *
   * @param key Key to compare against
   * @return std::size_t Number of elements with keys less than key
   *
   * Complexity: O(log n)
   */
  std::size_t order_of_key(const T &key) const {
    std::size_t res = 0;
    Node<T> *cur = root;
    while (cur) {
      if (comp(cur->key, key)) {
        res += get_subtree_size(cur->left) + 1;
        cur = cur->right;
      } else {
        cur = cur->left;
      }
    }
    return res;
  }
};

} // namespace avl_order_statistic_set

#endif // AVL_ORDER_STATISTIC_SET_HPP