Advertisement
Guest User

Untitled

a guest
Apr 25th, 2019
87
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 6.22 KB | None | 0 0
  1. #ifndef CPP_AVL_HPP
  2. #define CPP_AVL_HPP
  3.  
  4. #include <cstdint>
  5.  
  6. template<typename K>
  7. struct avl_node {
  8. using node_t = avl_node<K>;
  9.  
  10. node_t *cs_[2]{nullptr, nullptr};
  11. node_t *p_{nullptr};
  12. int8_t bf_{0};
  13. uint8_t c_idx_{0};
  14. K k_;
  15.  
  16. explicit avl_node(K k) : k_(k) {}
  17.  
  18. void set_child(int8_t c_idx, avl_node *child) {
  19. cs_[c_idx] = child;
  20. if (child != nullptr) {
  21. child->p_ = this;
  22. child->c_idx_ = c_idx;
  23. }
  24. }
  25.  
  26. void set_child_nonnull(int8_t c_idx, avl_node *child) {
  27. cs_[c_idx] = child;
  28. child->p_ = this;
  29. child->c_idx_ = c_idx;
  30. }
  31. };
  32.  
  33. template<typename K, typename K_OP>
  34. struct avl_tree {
  35. using node_t = avl_node<K>;
  36.  
  37. node_t *root_;
  38. size_t n_{0};
  39.  
  40. avl_tree() : root_(nullptr) {}
  41.  
  42. node_t *begin() const {
  43. node_t *prev = nullptr;
  44. for (node_t *node = root_; node != nullptr; node = node->cs_[0]) {
  45. prev = node;
  46. }
  47. return prev;
  48. }
  49.  
  50. node_t *end() const { return nullptr; }
  51.  
  52. node_t *rbegin() const {
  53. node_t *prev = nullptr;
  54. for (node_t *node = root_; node != nullptr; node = node->cs_[1]) {
  55. prev = node;
  56. }
  57. return prev;
  58. }
  59.  
  60. node_t *rend() const { return nullptr; }
  61.  
  62. node_t *next(node_t *node) const {
  63. if (node->cs_[1] != nullptr) {
  64. node = node->cs_[1];
  65. while (node->cs_[0] != nullptr) node = node->cs_[0];
  66. } else {
  67. for (;;) {
  68. uint8_t c_idx = node->c_idx_;
  69. node = node->p_;
  70. if (node == nullptr) return nullptr;
  71. if (c_idx == 0) break;
  72. }
  73. }
  74.  
  75. return node;
  76. }
  77.  
  78. node_t *prev(node_t *node) const {
  79. if (node->cs_[0] != nullptr) {
  80. node = node->cs_[0];
  81. while (node->cs_[1] != nullptr) node = node->cs_[1];
  82. } else {
  83. for (;;) {
  84. uint8_t c_idx = node->c_idx_;
  85. node = node->p_;
  86. if (node == nullptr) return nullptr;
  87. if (c_idx == 1) break;
  88. }
  89. }
  90.  
  91. return node;
  92. }
  93.  
  94. unsigned int rotate(node_t *node, node_t *parent, unsigned int reason) {
  95. // return 0 if h(node) does not change
  96. // return 1 if h(node) decrease by 1
  97. int l = reason;
  98. int r = 1 - l;
  99. auto *P = node;
  100. auto *C = P->cs_[l];
  101. auto *N = C->cs_[r];
  102. if (C->bf_ != 1 - 2 * l) {
  103. /* P
  104. * +-----+-----+
  105. * C
  106. * +--+--+
  107. * N
  108. * becomes:
  109. * C
  110. * +-----+-----+
  111. * P
  112. * +--+--+
  113. * N
  114. *
  115. * 1) l(0), P(-2), C( 0) -> P(-1), C(+1), return 0
  116. * 2) l(0), P(-2), C(-1) -> P( 0), C( 0), return 1
  117. * 1) l(1), P(+2), C( 0) -> P(+1), C(-1), return 0
  118. * 2) l(1), P(+2), C(+1) -> P( 0), C( 0), return 1
  119. */
  120. if (parent == nullptr) root_ = C;
  121. else parent->cs_[P->c_idx_] = C;
  122. C->p_ = parent;
  123. C->c_idx_ = P->c_idx_;
  124.  
  125. C->set_child_nonnull(r, P);
  126. P->set_child(l, N);
  127.  
  128. C->bf_ += 1 - 2 * l;
  129. P->bf_ = -C->bf_;
  130.  
  131. return C->bf_ == 0;
  132. } else {
  133. /* P
  134. * +-----+-----+
  135. * C
  136. * +--+--+
  137. * N
  138. * +-+-+
  139. * Nl Nr
  140. * becomes:
  141. * N
  142. * +-----+-----+
  143. * C P
  144. * +--+--+ +--+--+
  145. * Nl Nr
  146. *
  147. * 1) l(0), P(-2), C(+1), N(-1) -> P(+1), C( 0), N( 0), return 1
  148. * 2) l(0), P(-2), C(+1), N( 0) -> P( 0), C( 0), N( 0), return 1
  149. * 3) l(0), P(-2), C(+1), N(+1) -> P( 0), C(-1), N( 0), return 1
  150. * 4) l(1), P(+2), C(-1), N(+1) -> P(-1), C( 0), N( 0), return 1
  151. * 5) l(1), P(+2), C(-1), N( 0) -> P( 0), C( 0), N( 0), return 1
  152. * 6) l(1), P(+2), C(-1), N(-1) -> P( 0), C(+1), N( 0), return 1
  153. */
  154. auto *Nl = N->cs_[l];
  155. auto *Nr = N->cs_[r];
  156.  
  157. if (parent == nullptr) root_ = N;
  158. else parent->cs_[P->c_idx_] = N;
  159. N->p_ = parent;
  160. N->c_idx_ = P->c_idx_;
  161.  
  162. N->set_child_nonnull(l, C);
  163. N->set_child_nonnull(r, P);
  164. C->set_child(r, Nl);
  165. P->set_child(l, Nr);
  166.  
  167. P->bf_ = (N->bf_ == 2 * l - 1) ? -N->bf_ : 0;
  168. C->bf_ = (N->bf_ == 1 - 2 * l) ? -N->bf_ : 0;
  169. N->bf_ = 0;
  170.  
  171. return 1;
  172. }
  173. }
  174.  
  175. avl_node<K> *find(const K k) {
  176. if (root_ == nullptr) return nullptr;
  177. auto *cursor = root_;
  178. for (;;) {
  179. if (K_OP::eq(k, cursor->k_)) return cursor;
  180. int l = !K_OP::le(k, cursor->k_);
  181. if (cursor->cs_[l] != nullptr) {
  182. cursor = cursor->cs_[l];
  183. } else {
  184. return cursor;
  185. }
  186. }
  187. }
  188.  
  189. void insert(const K k) {
  190. n_++;
  191. auto *node = new node_t{k};
  192. if (root_ == nullptr) {
  193. root_ = node;
  194. } else {
  195. auto *cursor = root_;
  196. for (;;) {
  197. int l = !K_OP::le(k, cursor->k_);
  198. if (cursor->cs_[l] != nullptr) {
  199. cursor = cursor->cs_[l];
  200. } else {
  201. cursor->cs_[l] = node;
  202. node->p_ = cursor;
  203. node->c_idx_ = l;
  204. cursor->bf_ += 2 * l - 1;
  205. break;
  206. }
  207. }
  208.  
  209. while (!(cursor->bf_ == 0 || cursor->p_ == nullptr)) {
  210. if (cursor->p_->bf_ == 2 * cursor->c_idx_ - 1) {
  211. rotate(cursor->p_, cursor->p_->p_, cursor->c_idx_);
  212. break;
  213. }
  214. cursor->p_->bf_ += 2 * cursor->c_idx_ - 1;
  215. cursor = cursor->p_;
  216. }
  217. }
  218. }
  219.  
  220. void erase(K k) {
  221. n_--;
  222. auto *target = root_;
  223. while (!K_OP::eq(k, target->k_)) target = target->cs_[!K_OP::le(k, target->k_)];
  224.  
  225. node_t *to_delete = target;
  226. unsigned int l = (target->bf_ + 1) / 2;
  227. if (target->cs_[0] != nullptr && target->cs_[1] != nullptr) {
  228. int r = 1 - l;
  229. to_delete = target->cs_[l];
  230. while (to_delete->cs_[r] != nullptr) to_delete = to_delete->cs_[r];
  231. }
  232.  
  233. target->k_ = std::move(to_delete->k_);
  234.  
  235. if (to_delete->p_ == nullptr) {
  236. root_ = to_delete->cs_[l];
  237. if (root_ != nullptr) root_->p_ = nullptr;
  238. delete to_delete;
  239. return;
  240. }
  241. auto *cursor = to_delete->p_;
  242. auto c_idx = to_delete->c_idx_;
  243. cursor->set_child(c_idx, to_delete->cs_[l]);
  244. delete to_delete;
  245.  
  246. l = !c_idx;
  247. while (cursor) {
  248. auto *p = cursor->p_;
  249. auto idx = cursor->c_idx_;
  250. if (cursor->bf_ == 2 * l - 1) {
  251. if (!rotate(cursor, p, l)) return;
  252. } else {
  253. cursor->bf_ += 2 * l - 1;
  254. if (cursor->bf_ != 0) return;
  255. }
  256. cursor = p;
  257. l = !idx;
  258. }
  259. }
  260. };
  261.  
  262. #endif //CPP_AVL_HPP
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement