Advertisement
OIQ

Untitled

OIQ
Nov 7th, 2021
927
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.91 KB | None | 0 0
  1. #include <iostream>
  2. #include <string>
  3.  
  4. struct SimpleRBNode {
  5.     int value;
  6.     int left;
  7.     int right;
  8.     bool isBlack;
  9. };
  10.  
  11. class ArrayRBTree {
  12. public:
  13.     ArrayRBTree(int size, int root) {
  14.         array_tree_ = new SimpleRBNode[size + 1];
  15.         root_index_ = root;
  16.         array_tree_[0] = SimpleRBNode{0, -1, -1, true};
  17.         bst_ = new int[size];
  18.         index_ = 0;
  19.     }
  20.  
  21.     ~ArrayRBTree() {
  22.         delete[] array_tree_;
  23.         delete[] bst_;
  24.     }
  25.  
  26.     void add(const SimpleRBNode& rb_node, int index) {
  27.         array_tree_[index] = rb_node;
  28.     }
  29.  
  30.     bool checkRbTree() {
  31.         int result = dfsBlackCount(root_index_);
  32.         bool is_rb = dfs(root_index_, 0);
  33.         bool is_bst = true;
  34.         for (int i = 1; i < index_; ++i) {
  35.             if (bst_[i] <= bst_[i - 1]) {
  36.                 is_bst = false;
  37.                 break;
  38.             }
  39.         }
  40.         return array_tree_[root_index_].isBlack && is_rb && result != -1 && is_bst;
  41.     }
  42.  
  43. private:
  44.     bool dfs(int node, int parent) {
  45.         if (node == 0) {
  46.             return array_tree_[node].isBlack;
  47.         }
  48.  
  49.         if (!array_tree_[node].isBlack && !array_tree_[parent].isBlack) {
  50.             return false;
  51.         }
  52.  
  53.         bool left = dfs(array_tree_[node].left, node);
  54.         bst_[index_++] = array_tree_[node].value;
  55.         bool right = dfs(array_tree_[node].right, node);
  56.  
  57.         return left && right;
  58.     }
  59.  
  60.     int dfsBlackCount(int node) {
  61.         if (node == 0) {
  62.             return 0;
  63.         }
  64.         int count_left = dfsBlackCount(array_tree_[node].left);
  65.         int count_right = dfsBlackCount(array_tree_[node].right);
  66.         if (count_left >= 0 && count_right >= 0 && count_left == count_right) {
  67.             return count_left + (array_tree_[node].isBlack ? 1 : 0);
  68.         } else {
  69.             return -1;
  70.         }
  71.     }
  72.  
  73.     SimpleRBNode* array_tree_;
  74.     int root_index_;
  75.     int* bst_;
  76.     int index_;
  77. };
  78.  
  79. void parse(ArrayRBTree* tree) {
  80.     int index, value;
  81.     std::cin >> index >> value;
  82.  
  83.     std::string left_child, right_child, color;
  84.     std::cin >> left_child >> right_child >> color;
  85.  
  86.     SimpleRBNode node{value, left_child == "null" ? 0 : stoi(left_child),
  87.                       right_child == "null" ? 0 : stoi(right_child), color == "B"};
  88.  
  89.     tree->add(node, index);
  90. }
  91.  
  92. int execute(int size, int root) {
  93.     ArrayRBTree array_tree(size, root);
  94.  
  95.     for (int i = 0; i < size; ++i) {
  96.         parse(&array_tree);
  97.     }
  98.  
  99.     std::cout << (array_tree.checkRbTree() ? "YES" : "NO");
  100.  
  101.     return 0;
  102. }
  103.  
  104. void fastIO() {
  105.     std::ios_base::sync_with_stdio(false);
  106.     std::cin.tie(nullptr);
  107.     std::cout.tie(nullptr);
  108. }
  109.  
  110. int main(int argc, char** argv) {
  111.     fastIO();
  112.     int size, root;
  113.     std::cin >> size;
  114.     if (size == 0) {
  115.         std::cout << "YES";
  116.         return 0;
  117.     }
  118.     std::cin >> root;
  119.  
  120.     return execute(size, root);
  121. }
  122.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement