Seminar 11

#include <cassert>
#include <iostream>

struct AVLNode {
  int Key;
  int Height = 1;  // Высота дерева. 1 для листа.
  int Count = 1;  // Количество узлов в поддереве, включая сам узел.
  AVLNode* Left = nullptr;
  AVLNode* Right = nullptr;

  explicit AVLNode(int key) : Key(key) {}
  ~AVLNode() { delete Left; delete Right; }
};

int GetHeight(AVLNode* node) {
  return node ? node->Height : 0;
}

int GetCount(AVLNode* node) {
  return node ? node->Count : 0;
}

void UpdateHeight(AVLNode* node) {
  assert(node);
  node->Height = std::max(GetHeight(node->Left), GetHeight(node->Right)) + 1;
}

void UpdateCount(AVLNode* node) {
  assert(node);
  node->Count = GetCount(node->Left) + GetCount(node->Right) + 1;
}

int Balance(AVLNode* node) {
  assert(node);
  return GetHeight(node->Left) - GetHeight(node->Right);
}

class AVLTree {
 public:
  ~AVLTree() { delete root; }
  void Print() const { Print(root, 0); }
  int Add(int key);
  int Size() const { return GetCount(root); }
  int KStat(int k) const {
    assert(k >= 0 && k < GetCount(root));
    return KStat(k, root);
  }

 private:
  AVLNode* root = nullptr;

  static int Add(int key, AVLNode*& node);
  static void Print(AVLNode* node, int tabs);
  static void RotateLeft(AVLNode*& node);
  static void RotateRight(AVLNode*& node);
  static void FixBalance(AVLNode*& node);
  static int KStat(int k, AVLNode* node);
};

int AVLTree::Add(int key) {
  return Add(key, root);
}

// Тут можно вернуть статистику добавленного элемента
int AVLTree::Add(int key, AVLNode*& node) {
  if (!node) {
    node = new AVLNode(key);
    return 0;
  }
  int stat = 0;
  if (node->Key < key) {
    stat = Add(key, node->Right) + GetCount(node->Left) + 1;
  } else {
    stat = Add(key, node->Left);
  }
  // Восстановим высоту и баланс.
  UpdateHeight(node);
  UpdateCount(node);
  FixBalance(node);
  return stat;
}

void AVLTree::RotateLeft(AVLNode*& node) {
  AVLNode* right = node->Right;
  node->Right = right->Left;
  right->Left = node;
  node = right;

  UpdateHeight(node->Left);
  UpdateCount(node->Left);
  UpdateHeight(node);
  UpdateCount(node);
}

void AVLTree::RotateRight(AVLNode*& node) {
  AVLNode* left = node->Left;
  node->Left = left->Right;
  left->Right = node;
  node = left;

  UpdateHeight(node->Right);
  UpdateCount(node->Right);
  UpdateHeight(node);
  UpdateCount(node);
}

void AVLTree::FixBalance(AVLNode*& node) {
  int balance = Balance(node);
  if (balance == -2) {
    if (Balance(node->Right) == 1) {
      RotateRight(node->Right);
    }
    RotateLeft(node);
  } else if (balance == 2) {
    if (Balance(node->Left) == -1) {
      RotateLeft(node->Left);
    }
    RotateRight(node);
  }
}

int AVLTree::KStat(int k, AVLNode* node) {
  if (GetCount(node->Left) == k) {
    return node->Key;
  }
  if (GetCount(node->Left) > k) {
    return KStat(k, node->Left);
  }
  return KStat(k - GetCount(node->Left) - 1, node->Right);
}

void AVLTree::Print(AVLNode* node, int tabs) {
  if (node->Right) Print(node->Right, tabs + 1);
  for (int i = 0; i < tabs; ++i) std::cout << "\t";
  std::cout << node->Key << "(" << node->Count << ")" << std::endl;
  if (node->Left) Print(node->Left, tabs + 1);
}

// [.............3333333388888888................]
//               u       r

int main1() {
  AVLTree tree;
  tree.Add(5);
  tree.Add(6);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(8);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(9);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(14);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(7);
  tree.Print();
  std::cout << "===============\n";
  std::cout << tree.KStat(0) << std::endl;
  std::cout << tree.KStat(1) << std::endl;
  std::cout << tree.KStat(2) << std::endl;
  std::cout << tree.KStat(3) << std::endl;
  std::cout << tree.KStat(4) << std::endl;
  std::cout << tree.KStat(5) << std::endl;
  return 0;
}

struct Node {
  int Data;
  int Count = 1;  // Количество узлов в поддереве, включая сам узел.
  Node* Left = nullptr;
  Node* Right = nullptr;

  explicit Node(int data) : Data(data) {}
  ~Node() { delete Left; delete Right; }
};

int GetCount(Node* node) {
  return node ? node->Count : 0;
}
class TreeByImplicitKey {
 public:
  ~TreeByImplicitKey() { delete root; }
  void Print() const { Print(root, 0); }
  // Добавить узел data в позицию pos.
  void Add(int data, int pos) { Add(data, pos, root); }
  int Size() const { return GetCount(root); }
  int KStat(int k) const {
    assert(k >= 0 && k < GetCount(root));
    return KStat(k, root);
  }

 private:
  Node* root = nullptr;

  static void Add(int data, int pos, Node*& node);
  static void Print(Node* node, int tabs);
  static int KStat(int k, Node* node);
};

void TreeByImplicitKey::Add(int data, int pos, Node*& node) {
  if (!node) {
    node = new Node(data);
    return;
  }
  if (GetCount(node->Left) >= pos) {
    Add(data, pos, node->Left);
  } else {
    Add(data, pos - GetCount(node->Left) - 1, node->Right);
  }
}

void TreeByImplicitKey::Print(Node* node, int tabs) {
  if (node->Right) Print(node->Right, tabs + 1);
  for (int i = 0; i < tabs; ++i) std::cout << "\t";
  std::cout << node->Data << "(" << node->Count << ")" << std::endl;
  if (node->Left) Print(node->Left, tabs + 1);
}

int main() {
  TreeByImplicitKey tree;
  tree.Add(5, 0);
  tree.Add(6, 1);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(8, 0);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(9, 2);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(14, 3);
  tree.Print();
  std::cout << "===============\n";
  tree.Add(7, 1);
  tree.Print();
  std::cout << "===============\n";
  return 0;
}