Problem p02455 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdio>
#include <functional>
#include <vector>

#define fprintf(...) void(0)

template <class T> class RedBlackTree {
  struct Node {
    Node *left = nullptr, *right = nullptr, *parent = nullptr;
    T value;
    enum Color { Red = false, Black = true } color = Red;
    Node(T x) : value(x) {}
  } *root = nullptr;
  size_t size_ = 0;

  void left_rotate(Node *cur) {
    fprintf(stderr, "left_rotate(&%d)\\n", cur->value);
    Node *right = cur->right;
    assert(right != nullptr);
    cur->right = right->left;
    if (right->left != nullptr)
      right->left->parent = cur;
    right->parent = cur->parent;
    if (cur->parent == nullptr) {
      root = right;
    } else if (cur == cur->parent->left) {
      cur->parent->left = right;
    } else {
      cur->parent->right = right;
    }
    right->left = cur;
    cur->parent = right;
  }

  void right_rotate(Node *cur) {
    fprintf(stderr, "right_rotate(&%d)\\n", cur->value);
    Node *left = cur->left;
    assert(left != nullptr);
    cur->left = left->right;
    if (left->right != nullptr)
      left->right->parent = cur;
    left->parent = cur->parent;
    if (cur->parent == nullptr) {
      root = left;
    } else if (cur == cur->parent->right) {
      cur->parent->right = left;
    } else {
      cur->parent->left = left;
    }
    left->right = cur;
    cur->parent = left;
  }

  bool is_red(Node *node) const {
    // NIL is black
    return node != nullptr && node->color == Node::Red;
  }

public:
  RedBlackTree() = default;

  size_t size() const { return size_; }

  bool insert(T x) {
    if (root == nullptr) {
      root = new Node(x);
      root->color = Node::Black;
      ++size_;
      return true;
    }

    Node **child = &root, *parent = nullptr;
    while (*child != nullptr) {
      parent = *child;
      if ((*child)->value > x) {
        child = &(*child)->left;
      } else if ((*child)->value < x) {
        child = &(*child)->right;
      } else {
        return false;
      }
    }
    *child = new Node(x);
    ++size_;
    (*child)->parent = parent;
    Node *cur = *child;

    while (is_red(cur->parent)) {
      Node *gparent = cur->parent->parent;
      assert(gparent != nullptr);
      if (cur->parent == gparent->left) {
        Node *uncle = gparent->right;
        if (is_red(uncle)) {
          cur->parent->color = uncle->color = Node::Black;
          gparent->color = Node::Red;
          cur = gparent;
        } else {
          if (cur == cur->parent->right) {
            cur = cur->parent;
            left_rotate(cur);
          }
          cur->parent->color = Node::Black;
          cur->parent->parent->color = Node::Red;
          right_rotate(cur->parent->parent);
        }
      } else {
        Node *uncle = gparent->left;
        if (is_red(uncle)) {
          cur->parent->color = uncle->color = Node::Black;
          gparent->color = Node::Red;
          cur = gparent;
        } else {
          if (cur == cur->parent->left) {
            cur = cur->parent;
            right_rotate(cur);
          }
          cur->parent->color = Node::Black;
          cur->parent->parent->color = Node::Red;
          left_rotate(cur->parent->parent);
        }
      }
      fprintf(stderr, "cur: %d\\n", cur->value);
    }
    root->color = Node::Black;
    return true;
  }

  bool contains(T x) const {
    Node *cur = root;
    while (cur != nullptr) {
      if (cur->value < x) {
        cur = cur->right;
      } else if (cur->value > x) {
        cur = cur->left;
      } else {
        return true;
      }
    }
    return false;
  }

  bool verify() const {
    if (root == nullptr)
      return true;

    bool violated = false;
    if (root->color == Node::Red) {
      violated = true;
      fprintf(stderr, "root node is red\\n");
    }

    bool reached = false;
    size_t bheight = -1;
    std::function<void(Node *, int, size_t)> dfs = [&](Node *subroot, int depth,
                                                       size_t bh) {
      if (subroot->color == Node::Red) {
        if (subroot->parent && subroot->parent->color == Node::Red) {
          violated = true;
          fprintf(stderr, "two red nodes are connected\\n");
        }
      } else {
        ++bh;
      }

      if (subroot->left == nullptr || subroot->right == nullptr) {
        if (!reached) {
          bheight = bh;
        } else if (bheight != bh) {
          violated = true;
          fprintf(stderr,
                  "different number of black nodes in two root-leaf paths\\n");
        }
      }

      if (subroot->right != nullptr)
        dfs(subroot->right, depth + 1, bh);

      fprintf(stderr, "%s%*s%d%s\\n",
              (subroot->color == Node::Red) ? "\\x1b[31;1m" : "\\x1b[37;1m",
              depth * 2 + 2, "  ", subroot->value, "\\x1b[0m");

      if (subroot->left != nullptr)
        dfs(subroot->left, depth + 1, bh);
    };
    dfs(root, 0, 0);
    return !violated;
  }
};

int main() {
  RedBlackTree<int> rbt;

  int q;
  scanf("%d", &q);
  for (int i = 0; i < q; ++i) {
    int t, x;
    scanf("%d", &t);
    if (t == 0) {
      scanf("%d", &x);
      rbt.insert(x);
      printf("%zu\\n", rbt.size());
    } else if (t == 1) {
      scanf("%d", &x);
      printf("%d\\n", rbt.contains(x));
    }
  }
  assert(rbt.verify());
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdio>
#include <functional>
#include <vector>

#define fprintf(...) void(0)

template <class T> class RedBlackTree {
  struct Node {
    Node *left = nullptr, *right = nullptr, *parent = nullptr;
    T value;
    enum Color { Red = false, Black = true } color = Red;
    Node(T x) : value(x) {}
  } *root = nullptr;
  size_t size_ = 0;

  void left_rotate(Node *cur) {
    fprintf(stderr, "left_rotate(&%d)\\n", cur->value);
    Node *right = cur->right;
    assert(right != nullptr);
    cur->right = right->left;
    if (right->left != nullptr)
      right->left->parent = cur;
    right->parent = cur->parent;
    if (cur->parent == nullptr) {
      root = right;
    } else if (cur == cur->parent->left) {
      cur->parent->left = right;
    } else {
      cur->parent->right = right;
    }
    right->left = cur;
    cur->parent = right;
  }

  void right_rotate(Node *cur) {
    fprintf(stderr, "right_rotate(&%d)\\n", cur->value);
    Node *left = cur->left;
    assert(left != nullptr);
    cur->left = left->right;
    if (left->right != nullptr)
      left->right->parent = cur;
    left->parent = cur->parent;
    if (cur->parent == nullptr) {
      root = left;
    } else if (cur == cur->parent->right) {
      cur->parent->right = left;
    } else {
      cur->parent->left = left;
    }
    left->right = cur;
    cur->parent = left;
  }

  bool is_red(Node *node) const {
    // NIL is black
    return node != nullptr && node->color == Node::Red;
  }

public:
  RedBlackTree() = default;

  size_t size() const { return size_; }

  bool insert(T x) {
    if (root == nullptr) {
      root = new Node(x);
      root->color = Node::Black;
      ++size_;
      return true;
    }

    Node **child = &root, *parent = nullptr;
    while (*child != nullptr) {
      parent = *child;
      if ((*child)->value > x) {
        child = &(*child)->left;
      } else if ((*child)->value < x) {
        child = &(*child)->right;
      } else {
        return false;
      }
    }
    *child = new Node(x);
    ++size_;
    (*child)->parent = parent;
    Node *cur = *child;

    while (is_red(cur->parent)) {
      Node *gparent = cur->parent->parent;
      assert(gparent != nullptr);
      if (cur->parent == gparent->left) {
        Node *uncle = gparent->right;
        if (is_red(uncle)) {
          cur->parent->color = uncle->color = Node::Black;
          gparent->color = Node::Red;
          cur = gparent;
        } else {
          if (cur == cur->parent->right) {
            cur = cur->parent;
            left_rotate(cur);
          }
          cur->parent->color = Node::Black;
          cur->parent->parent->color = Node::Red;
          right_rotate(cur->parent->parent);
        }
      } else {
        Node *uncle = gparent->left;
        if (is_red(uncle)) {
          cur->parent->color = uncle->color = Node::Black;
          gparent->color = Node::Red;
          cur = gparent;
        } else {
          if (cur == cur->parent->left) {
            cur = cur->parent;
            right_rotate(cur);
          }
          cur->parent->color = Node::Black;
          cur->parent->parent->color = Node::Red;
          left_rotate(cur->parent->parent);
        }
      }
      fprintf(stderr, "cur: %d\\n", cur->value);
    }
    root->color = Node::Black;
    return true;
  }

  bool contains(T x) const {
    Node *cur = root;
    while (cur != nullptr) {
      if (cur->value < x) {
        cur = cur->right;
      } else if (cur->value > x) {
        cur = cur->left;
      } else {
        return true;
      }
    }
    return false;
  }

  bool verify() const {
    if (root == nullptr)
      return true;

    bool violated = false;
    if (root->color == Node::Red) {
      violated = true;
      fprintf(stderr, "root node is red\\n");
    }

    bool reached = false;
    size_t bheight = -1;
    std::function<void(Node *, int, size_t)> dfs = [&](Node *subroot, int depth,
                                                       size_t bh) {
      if (subroot->color == Node::Red) {
        if (subroot->parent && subroot->parent->color == Node::Red) {
          violated = true;
          fprintf(stderr, "two red nodes are connected\\n");
        }
      } else {
        ++bh;
      }

      if (subroot->left == nullptr || subroot->right == nullptr) {
        if (!reached) {
          bheight = bh;
        } else if (bheight != bh) {
          violated = true;
          fprintf(stderr,
                  "different number of black nodes in two root-leaf paths\\n");
        }
      }

      if (subroot->right != nullptr)
        dfs(subroot->right, depth + 1, bh);

      fprintf(stderr, "%s%*s%d%s\\n",
              (subroot->color == Node::Red) ? "\\x1b[31;1m" : "\\x1b[37;1m",
              depth * 2 + 2, "  ", subroot->value, "\\x1b[0m");

      if (subroot->left != nullptr)
        dfs(subroot->left, depth + 1, bh);
    };
    dfs(root, 0, 0);
    return !violated;
  }
};

int main() {
  RedBlackTree<int> rbt;

  int q;
  scanf("%d", &q);
  for (int i = 0; i < q; ++i) {
    int t, x;
    scanf("%d", &t);
    if (t == 0) {
      scanf("%d", &x);
      rbt.insert(x);
      printf("%zu\\n", rbt.size());
      assert(rbt.verify());
    } else if (t == 1) {
      scanf("%d", &x);
      printf("%d\\n", rbt.contains(x));
    }
  }
}