TreeMap
完整代码
const RED = "0";
const BLACK = "1";
class Node {
left;
right;
parent;
constructor(key, value, color) {
this.key = key;
this.value = value;
this.color = color ?? RED;
}
compareTo(target) {
let targetKey = target;
if (target instanceof Node) {
targetKey = target.key;
}
return this.key - targetKey;
}
}
class TreeSet {
constructor() {
this.size = 0;
this.root = null;
this.dummyNode = new Node(null, null, BLACK);
}
leftRotate(node) {
let right = node.right;
node.right = right.left;
if (right.left !== this.dummyNode) right.left.parent = node;
right.parent = node.parent;
if (node.parent === this.dummyNode) this.root = right;
else if (node.parent.left == node) node.parent.left = right;
else node.parent.right = right;
right.left = node;
node.parent = right;
}
rightRotate(node) {
let left = node.left;
node.left = left.right;
if (left.right !== this.dummyNode) left.right.parent = node;
left.parent = node.parent;
if (node.parent === this.dummyNode) {
this.root = left;
} else if (node.parent.left == node) {
node.parent.left = left;
} else {
node.parent.right = left;
}
left.right = node;
node.parent = left;
}
transplant(aNode, bNode) {
if (aNode.parent === this.dummyNode) {
this.root = bNode;
} else if (aNode.parent.left === aNode) {
aNode.parent.left = bNode;
} else {
aNode.parent.right = bNode;
}
bNode.parent = aNode.parent;
}
fixAdd(node) {
while (node.parent.color == RED) {
if (node.parent == node.parent.parent.left) {
let uncle = node.parent.parent.right;
if (uncle && uncle.color == RED) {
uncle.color = BLACK;
node.parent.color = BLACK;
node.parent.parent.color = RED;
node = node.parent.parent;
} else if (node == node.parent.right) {
node = node.parent;
this.leftRotate(node);
} else {
node.parent.color = BLACK;
node.parent.parent.color = RED;
this.rightRotate(node.parent.parent);
}
} else {
let uncle = node.parent.parent.left;
if (uncle && uncle.color == RED) {
uncle.color = BLACK;
node.parent.color = BLACK;
node.parent.parent.color = RED;
node = node.parent.parent;
} else if (node == node.parent.left) {
node = node.parent;
this.rightRotate(node);
} else {
node.parent.color = BLACK;
node.parent.parent.color = RED;
this.leftRotate(node.parent.parent);
}
}
}
this.root.color = BLACK;
}
fixRemove(node) {
while (node != this.root && node.color == BLACK) {
if (node === node.parent.left) {
let brother = node.parent.right;
if (brother.color == RED) {
brother.color = BLACK;
node.parent.color = RED;
this.leftRotate(node.parent);
brother = node.parent.right;
}
if (brother.left.color === BLACK && brother.right.color === BLACK) {
brother.color = RED;
node = node.parent;
} else if (brother.right.color === BLACK) {
brother.left.color = BLACK;
brother.color = RED;
this.rightRotate(brother);
brother = node.parent.right;
} else {
brother.color = node.parent.color;
node.parent.color = BLACK;
brother.right.color = BLACK;
this.leftRotate(node.parent);
node = this.root;
}
} else {
let brother = node.parent.left;
if (brother.color == RED) {
brother.color = BLACK;
node.parent.color = RED;
this.rightRotate(node.parent);
brother = node.parent.left;
}
if (brother.right.color === BLACK && brother.left.color === BLACK) {
brother.color = RED;
node = node.parent;
} else if (brother.left.color === BLACK) {
brother.right.color = BLACK;
brother.color = RED;
this.leftRotate(brother);
brother = node.parent.left;
} else {
brother.color = node.parent.color;
node.parent.color = BLACK;
brother.left.color = BLACK;
this.rightRotate(node.parent);
node = this.root;
}
}
}
if (node) {
node.color = BLACK;
}
}
add(key, value) {
let parent = this.dummyNode;
let current = this.root;
this.size++;
while (current !== null && current !== this.dummyNode) {
parent = current;
if (current.compareTo(key) > 0) {
current = current.left;
} else if (current.compareTo(key) < 0) {
current = current.right;
} else {
current.value = value;
this.size--;
return;
}
}
const node = new Node(key, value);
node.parent = parent;
if (parent === this.dummyNode) {
this.root = node;
} else if (node.compareTo(parent) < 0) {
parent.left = node;
} else {
parent.right = node;
}
node.left = this.dummyNode;
node.right = this.dummyNode;
this.fixAdd(node);
}
remove(key) {
let node = this.getNode(key);
if (node) {
return this.removeNode(node);
}
return false;
}
removeNode(node) {
let remove = node;
let removeColor = remove.color;
let replace;
if (node.left === this.dummyNode) {
replace = node.right;
this.transplant(node, node.right);
} else if (node.right === this.dummyNode) {
replace = node.left;
this.transplant(node, node.left);
} else {
remove = this.getSucceed(node);
removeColor = remove.color;
replace = remove.right;
if (remove.parent == node) {
replace.parent = remove;
} else {
this.transplant(remove, remove.right);
remove.right = node.right;
remove.right.parent = remove;
}
this.transplant(node, remove);
remove.left = node.left;
remove.left.parent = remove;
remove.color = node.color;
}
if (removeColor === BLACK) {
this.fixRemove(replace);
}
this.size--;
return true;
}
get(key) {
return this.getNode(key);
}
first() {
let current = this.root;
while (current.left !== this.dummyNode) {
current = current.left;
}
return current;
}
getNode(key) {
let current = this.root;
while (current !== null && current !== this.dummyNode) {
if (current.compareTo(key) > 0) {
current = current.left;
} else if (current.compareTo(key) < 0) {
current = current.right;
} else {
return current;
}
}
return current;
}
getSucceed(node) {
let right = node.right;
let succeed = right;
while (right !== this.dummyNode) {
succeed = right;
right = right.left;
}
return succeed;
}
levelRraversal() {
const result = [];
const queue = [];
queue.push(this.root);
while (queue.length) {
let temp = queue.shift();
result.push(temp);
if (temp.left !== this.dummyNode) queue.push(temp.left);
if (temp.right !== this.dummyNode) queue.push(temp.right);
}
return result;
}
}
测试用例
const treeSet = new TreeSet();
const array = [
12, 1, 9, 2, 0, 11, 7, 19, 4, 15, 18, 5, 14, 13, 10, 16, 6, 3, 8, 17,
];
array.forEach((item) => {
treeSet.add(item, item);
});
console.log(treeSet.levelRraversal())
treeSet.remove(12);
treeSet.remove(1);
treeSet.remove(9);
treeSet.remove(2);
treeSet.remove(0);
treeSet.remove(11);
treeSet.remove(7);