跳到主要内容

TreeSet

完整代码


class Node {
left;
right;
parent;
constructor(key) {
this.key = key;
}
compareTo(target) {
let targetKey = target;
if (target instanceof Node) {
targetKey = target.key;
}
return this.key - targetKey.key;
}
}
class TreeSet {
constructor() {
this.size = 0;
this.root = null;
}
transplant(aNode, bNode) {
if (!aNode.parent) {
this.root = bNode;
} else if (aNode.parent.left === aNode) {
aNode.parent.left = bNode;
} else {
aNode.parent.right = bNode;
}
if (bNode) {
bNode.parent = aNode.parent;
}
}
add(key) {
let parent = null;
let current = this.root;
this.size++;
while (current) {
parent = current;
if (current.compareTo(key) > 0) {
current = current.left;
} else if (current.compareTo(key) < 0) {
current = current.right;
} else {
current.key = key;
this.size--;
return;
}
}
const node = new Node(key);
if (!parent) {
this.root = node;
} else {
if (node.compareTo(parent) < 0) {
parent.left = node;
} else {
parent.right = node;
}
node.parent = parent;
}
}
remove(key) {
let node = this.getNode(key);
if (node) {
return this.removeNode(node);
}
return false;
}
removeNode(node) {
let remove = node;
let replace;
if (!node.left) {
replace = node.right;
this.transplant(node, node.right);
} else if (!node.right) {
replace = node.left;
this.transplant(node, node.left);
} else {
remove = this.getSucceed(node);
replace = remove.right;
if (remove.parent == node) {
if (replace) {
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;
}
this.size--;
return true;
}
get(key) {
const node = this.getNode(key);
return node;
}
getNode(key) {
let current = this.root;
while (current) {
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) {
succeed = right;
right = right.left;
}
return succeed;
}
}

测试用例