fork download
  1. #include <iostream>
  2.  
  3. enum Colour {
  4. Red = 0,
  5. Black = 1
  6. };
  7.  
  8. class Node {
  9. public:
  10. int value;
  11. Colour colour = Red;
  12. int num_left = 0;
  13. int num_right = 0;
  14. int num_bigger = 0;
  15. Node * parent = NULL;
  16. Node * left = NULL;
  17. Node * right = NULL;
  18. public:
  19. Node(int val) : value(val) { }
  20. Node(int val, Node * NIL) : value(val), parent(NIL), left(NIL), right(NIL) { }
  21. };
  22.  
  23. class RB_TREE {
  24. private:
  25. Node * root;
  26. Node * NIL;
  27. int cnt = 0;
  28. public:
  29. RB_TREE();
  30. void Insert(int val);
  31. void Insert_Fix(Node * node);
  32. void Left_Rotate(Node * node);
  33. void Right_Rotate(Node * node);
  34. int getCount() { return cnt; }
  35. };
  36.  
  37. RB_TREE::RB_TREE() {
  38. NIL = new Node(-1);
  39. NIL->colour = Black;
  40. root = NIL;
  41. }
  42.  
  43. void RB_TREE::Insert(int val) {
  44. Node * node = new Node(val, NIL);
  45.  
  46. Node * pos = NIL;
  47. Node * current = root;
  48. //std::cout << "--- Node: " << val << " ---\n"; //temp
  49.  
  50. while(current != NIL) {
  51. pos = current;
  52. //std::cout << "pos: " << pos->value << '\n'; //temp
  53.  
  54. if(node->value < current->value) {
  55. current->num_left++;
  56.  
  57. current = current->left;
  58. if(current !=NIL) {
  59. //std::cout << "[LEFT]Node " << current->value
  60. // << "'s num_bigger Before: " << current->num_bigger << " -> "; //temp
  61. current->num_bigger =
  62. current->parent->num_bigger + current->num_right + 1;
  63. //std::cout << "After: " << current->num_bigger << '\n'; //temp
  64. }
  65. }
  66. else {
  67. if(current->right != NIL) {
  68. //std::cout << "{RIGHT}Node " << current->right->value
  69. // << "'s num_bigger Before: " << current->right->num_bigger << " -> "; // temp
  70. current->right->num_bigger = current->num_bigger - current->right->num_left - 1;
  71.  
  72. //std::cout << "After: " << current->right->num_bigger << '\n'; //temp
  73. }
  74.  
  75. //std::cout << "[RIGHT]Node " << current->value
  76. // << "'s num_bigger Before: " << current->num_bigger << " -> "; // temp
  77. current->num_bigger++;
  78. current->num_right++;
  79.  
  80. //std::cout << "After: " << current->num_bigger << '\n'; //temp
  81.  
  82. current = current->right;
  83. }
  84. }
  85.  
  86. node->parent = pos;
  87.  
  88. if(pos == NIL) {
  89. root = node;
  90. }
  91. else if(node->value < pos->value) {
  92. pos->left = node;
  93. node->num_bigger = pos->num_bigger + 1;
  94. }
  95. else {
  96. pos->right = node;
  97. node->num_bigger = pos->num_bigger - 1;
  98. }
  99.  
  100. cnt += node->num_bigger;
  101. //std::cout << "Final num_bigger: " << node->num_bigger << '\n'; // temp
  102.  
  103. Insert_Fix(node);
  104. }
  105.  
  106. void RB_TREE::Insert_Fix(Node * node) {
  107. while(node->parent->colour == Red) {
  108. if(node->parent == node->parent->parent->left) {
  109. Node * uncle = node->parent->parent->right;
  110.  
  111. if(uncle->colour == Red) {
  112. node->parent->colour = Black;
  113. uncle->colour = Black;
  114.  
  115. node->parent->parent->colour = Red;
  116. node = node->parent->parent;
  117. }
  118. else{
  119. if(node == node->parent->right) {
  120. node = node->parent;
  121. Left_Rotate(node);
  122. }
  123.  
  124. node->parent->colour = Black;
  125. node->parent->parent->colour = Red;
  126. Right_Rotate(node->parent->parent);
  127. }
  128. }
  129. else {
  130. Node * uncle = node->parent->parent->left;
  131.  
  132. if(uncle->colour == Red) {
  133. node->parent->colour = Black;
  134. uncle->colour = Black;
  135.  
  136. node->parent->parent->colour = Red;
  137. node = node->parent->parent;
  138. }
  139. else{
  140. if(node == node->parent->left) {
  141. node = node->parent;
  142. Right_Rotate(node);
  143. }
  144.  
  145. node->parent->colour = Black;
  146. node->parent->parent->colour = Red;
  147. Left_Rotate(node->parent->parent);
  148. }
  149. }
  150. }
  151.  
  152. root->colour = Black;
  153. }
  154.  
  155. void RB_TREE::Left_Rotate(Node * node) {
  156. Node * other = node->right;
  157. node->right = other->left;
  158.  
  159. node->num_right -= other->num_right + 1;
  160. other->num_left += node->num_left + 1;
  161.  
  162. if(other->left != NIL) {
  163. other->left->parent = node;
  164. }
  165.  
  166. other->parent = node->parent;
  167.  
  168. if(node->parent == NIL) {
  169. root = other;
  170. }
  171. else if(node == node->parent->left) {
  172. node->parent->left = other;
  173. }
  174. else {
  175. node->parent->right = other;
  176. }
  177.  
  178. other->left = node;
  179. node->parent = other;
  180. }
  181.  
  182. void RB_TREE::Right_Rotate(Node * node) {
  183. Node * other = node->left;
  184. node->left = other->right;
  185.  
  186. node->num_left -= other->num_left + 1;
  187. other->num_right += other->num_right + 1;
  188.  
  189. if(other->right != NIL) {
  190. other->right->parent = node;
  191. }
  192.  
  193. other->parent = node->parent;
  194.  
  195. if(node->parent == NIL) {
  196. root = other;
  197. }
  198. else if(node == node->parent->right) {
  199. node->parent->right = other;
  200. }
  201. else {
  202. node->parent->left = other;
  203. }
  204.  
  205. other->right = node;
  206. node->parent = other;
  207. }
  208.  
  209. int main() {
  210. RB_TREE rbt;
  211. int num, val;
  212.  
  213. std::cin >> num;
  214.  
  215. for(int i = 0; i < num; i++) {
  216. std::cin >> val;
  217. rbt.Insert(val);
  218. }
  219.  
  220. std::cout << rbt.getCount() << '\n';
  221.  
  222. return 0;
  223. }
  224.  
Success #stdin #stdout 0.01s 5572KB
stdin
24
63 16 24 7 29 57 65 26 36 32 50 5 34 1 18 15 49 9 47 53 10 35 76 79
stdout
122