Красно-чёрное дерево

Материал из Олимпиадное программирование в УлГТУ
Перейти к навигации Перейти к поиску

Правила

  • Каждый узел является красным или чёрным.
  • Корень и (воображаемые) nullptr-потомки листьев — чёрные.
  • Родитель красного узла обязан быть чёрным (потомок и родитель не могут быть красными одновременно, не может быть двух красных узлов подряд);
  • Все простые пути пути от любой вершины до любого листа в её поддереве содержат одинаковое количество чёрных узлов.

Код

class RBTree {
    struct Node {
        int key;
        Node *l = nullptr, *r = nullptr, *p = nullptr;
        char color = 'R';

        Node *other(Node *n) {
            return l == n ? r : l;
        }
    } *root = nullptr;

    bool find(Node *node, int key) const {
        if (!node)
            return false;
        if (key == node->key)
            return true;
        return find(key < node->key ? node->l : node->r, key);
    }

    void rotateLeft(Node *x) {
        Node *y = x->r;
        x->r = y->l;
        if (y->l)
            y->l->p = x;
        y->p = x->p;
        if (!x->p)
            root = y;
        else if (x == x->p->l)
            x->p->l = y;
        else
            x->p->r = y;
        y->l = x;
        x->p = y;
    }

    void rotateRight(Node *x) {
        Node *y = x->l;
        x->l = y->r;
        if (y->r)
            y->r->p = x;
        y->p = x->p;
        if (!x->p)
            root = y;
        else if (x == x->p->r)
            x->p->r = y;
        else
            x->p->l = y;
        y->r = x;
        x->p = y;
    }

    void fix(Node *me) {
        if (me->color == 'B')
            return;

        Node *dad = me->p;
        if (dad && dad->color == 'B')
            return;

        Node *grandDad = dad->p, *uncle = grandDad ? grandDad->other(dad) : nullptr;
        if (uncle && uncle->color == 'R') {
            dad->color = uncle->color = 'B';
            grandDad->color = 'R';
            return;
        }

        if (me == dad->r && dad == grandDad->l) {
            rotateLeft(dad);
            return fix(me->l);
        }

        if (me == dad->l && dad == grandDad->r) {
            rotateRight(dad);
            return fix(me->r);
        }

        dad->color = 'B';
        grandDad->color = 'R';
        if (me == dad->l) {
            rotateRight(grandDad);
        } else {
            rotateLeft(grandDad);
        }
    }

    void insert(Node *&node, int key) {
        if (!node) {
            node = new Node({ key });
            return;
        }
        if (key < node->key) {
            insert(node->l, key);
            node->l->p = node;
            fix(node->l);
        } else {
            insert(node->r, key);
            node->r->p = node;
            fix(node->r);
        }
    }

public:
    bool find(int key) const {
        return find(root, key);
    }

    void insert(int key) {
        insert(root, key);
        root->color = 'B';
    }
};

Ссылки

Теория:

Демонстрация:

Код: