数据结构 - 红黑树

---

介绍

红黑树可以在 $O(logN)$ 时间内完成查找，插入，删除操作。相比 AVL ，牺牲了部分平衡性以在插入、删除操作时减少旋转操作，整体性能优于 AVL。 C++ STL 的 map 就是用红黑树实现的。

满足以下 5 个条件：

1. 每个节点非黑即红
2. 视为树（包含部分子树）的根节点是黑色的
3. 叶节点 （NIL） 是黑色的
4. 如果一个节点是红色，则它的两个子节点是黑色的
5. 从根节点出发到所有叶节点的路径上，黑色节点数量相同

4、5 两条规则使得红黑树从根节点到每个叶节点 （NIL）的最长路径最多是最短路径的两倍

编码细节

#include <stdio.h>
#include <stdlib.h>

typedef struct Node {
	int key;
	int color; // 0 red 1 black 2 double black
	struct Node *lchild, *rchild;
} Node;

#define RED 0
#define BLACK 1
#define DBLACK 2

// 所有的叶子节点都指向 NIL
Node __NIL;
#define NIL (&__NIL)
__attribute__((constructor))
void init_NIL() {
	NIL->key = 0;
	NIL->color = BLACK; // 叶子节点默认是 black 
	NIL->lchild = NIL->rchild = NIL;
	return ;
}

Node *getNewNode(int key) {
	Node *p = (Node *) malloc(sizeof(Node));
	p->key = key;
	p->color = RED; // 叶子节点创建时 默认是红色，使得路径上黑色节点数量保持不变
	p->lchild = p->rchild = NIL; // 叶子节点都指向 p
	return p;
}

int has_red_child(Node *root) {
	return root->lchild->color == RED || root->rchild->color == RED;
}

Node *left_rotate(Node *root) {
	Node *temp = root->rchild; // 左旋，右子节点作为父节点
	root->rchild = temp->lchild;
	temp->lchild = root;
	return temp;
}

Node *right_rotate(Node *root) {
	Node *temp = root->lchild;
	root->lchild = temp->rchild;
	temp->rchild = root;
	return temp;
}

Node *insert_maintain(Node *root) {
	// 子节点不存在红色的不需要调整
	if (!has_red_child(root)) return root;
	int flag = 0;
	// 如果子节点两个都是红色的
	if (root->lchild->color == RED && root->rchild->color == RED) {
		goto insert_end;
	}
	// 需要从祖父节点看父辈节点，孙辈节点关系
	// 出现两个红色节点临近连接，不满足规则4
	if (root->lchild->color == RED && has_red_child(root->lchild)) flag = 1;
	if (root->rchild->color == RED && has_red_child(root->rchild)) flag = 2;
	// 不需要调整
	if (flag == 0) return root;
	// 旋转为满足第四条规则
	if (flag == 1) {
		// LR 情况
		if (root->lchild->rchild->color == RED) {
			root->lchild = left_rotate(root->lchild);
		}
		root = right_rotate(root);
	} else {
		// RL 情况
		if (root->rchild->lchild->color == RED) {
			root->rchild = right_rotate(root->rchild);
		}
		root = left_rotate(root);
	}
insert_end:
	// 此处代码红色是上浮的，实际效果同常见的红黑树实现是一样的，黑色节点路径数不变
	// 变换策略：子节点两个都是红色的，子节点红色上浮到父节点，子节点变黑色，变换后路径上黑色节点数量不变
	// 变换策略：叔父节点和一个子节点颜色不同，旋转后此处代码也实现了旋转后上下节点颜色交换
	root->color = RED;
	root->lchild->color = root->rchild = BLACK;
	return root;
}

Node *__insert(Node *root, int key) {
	if (root == NIL) return getNewNode(key);
	if (root->key == key) return root;
	if (key < root->key) {
		root->lchild = __insert(root->lchild, key);
	} else {
		root->rchild = __insert(root->rchild, key);
	}
	// 插入后从子节点回溯调整
	return insert_maintain(root);
}

Node *insert(Node *root, int key) {
	root = __insert(root, key);
	root->color = BLACK; // 使得整棵树最终根节点始终为黑色，抵消 __insert 红色上浮结果
	return root;
}

// 找前驱节点
Node *predecessor(Node *root) {
	Node *temp = root->lchild;
	while (temp->rchild != NIL) temp = temp->rchild;
	return temp;
}

// 删除红色节点，不会对红黑树的平衡产生影响
// 删除度为1的黑色节点，会产生唯一子孩子，一定是红色，不会产生删除调整
// 删除度为0的黑色节点，会产生一个双重黑的 NIL 节点，就是为了干掉双重黑
Node *erase_maintain(Node *root) {
	// 如果不存在子节点需要路径不一致处理返回
	if (root->lchild->color != DBLACK && root->rchild->color != DBLACK) return root;
	if (has_red_child(root)) {
		int flag = 0;
		// 使得有红色子节点的节点旋转，用红色节点抵消
		root->color = RED;
		if (root->lchild->color == RED) {
			root = right_rotate(root); flag = 1;
		} else {
			root = left_rotate(root); flag = 2;
		}
		root->color = BLACK;
		if (flag == 1) root->rchild = erase_maintain(root->rchild);
		else root->lchild = erase_maintain(root->lchild);
		return root;
	}
	// 情况 1： 双重黑节点的兄弟节点有两个黑色节点，兄弟节点和双重黑节点 - 1黑，双重黑情况向上传毒
	if ((root->lchild->color == DBLACK && !has_red_child(root->rchild)) || 
		(root->rchild->color == DBLACK && !has_red_child(root->lchild))) {
		root->lchild->color -= 1;
		root->rchild->color -= 1;
		root->color += 1;
		return root; // 若本身为红色没事，若为黑色，double black 向上传递
	}
	// 情况 2：双重黑兄弟节点是黑色，且兄弟节点中有红色子节点，可以调整
	// 兄弟节点 RR，左旋，把新根为原根的颜色，新根两个子节点为黑色，抵消双黑数量影响
	// RL，先小右旋，对调颜色
	if (root->lchild->color == DBLACK) {
		// 必须使用 != 判断，不用 == BLACK 是因为 NIL 也存在 color 属性
		if (root->rchild->rchild->color != RED) {
			root->rchild->color = RED;
			root->rchild = right_rotate(root->rchild);
			root->rchild->color = BLACK;
		}
		root = left_rotate(root);
		root->color = root->lchild->color;
	} else {
		if (root->lchild->lchild != RED) {
			root->lchild->color = RED;
			root->lchild = left_rotate(root->lchild);
			root->lchild->color = BLACK;
		}
		root = right_rotate(root);
		root->color = root->rchild->color;
	}
	root->lchild->color = root->rchild->color = BLACK;
	return root;
}

Node *__erase(Node *root, int key) {
	if (root == NIL) return NIL;
	if (key < root->key) {
		root->lchild = __erase(root->lchild, key);
	} else if (key > root->key) {
		root->rchild = __erase(root->rchild, key);
	} else if (root->lchild == NIL || root->rchild == NIL) {
		Node *temp = root->lchild != NIL ? root->lchild : root->rchild;
		temp->color += root->color; // 被删去的节点使用左节点或右节点顶替
		// 如果删去的是红色没关系，删去的是黑色，要加上删去黑色信息，所以有 double black 颜色，告知路径调整
		free(root);
		return temp;
	} else {
		Node *temp = predecessor(root);
		root->key = temp->key;
		root->lchild = __erase(root->lchild, temp->key);
	}
	return erase_maintain(root);
}

Node *erase(Node *root, int key) {
	root = __erase(root, key);
	root->color = BLACK;
	return root;
}

void clear(Node *root) {
	if (root == NIL) return ;
	clear(root->lchild);
	clear(root->rchild);
	free(root);
	return ;
}

void print(Node *root) {
	printf("(%d | %d, %d, %d)\n",
		root->color, root->key,
		root->lchild->key,
		root->rchild->key
	);
	return ;
}

void output(Node *root) {
	if (root == NIL) return ;
	print(root);
	output(root->lchild);
	output(root->rchild);
	return ;
}

int main() {
	int op, val;
	Node *root = NIL;
	while (~scanf("%d%d", &op, &val)) {
		switch (op) {
			case 1: root = insert(root, val); break;
			case 2: root = erase(root, val); break;
		}
		output(root);
		printf("-------------------\n");
	}
	return 0;
}