mirror of
https://github.com/irislib/iris-messenger.git
synced 2024-10-18 14:13:21 +00:00
336 lines
9.1 KiB
TypeScript
336 lines
9.1 KiB
TypeScript
// just for comparison with @/utils/SortedMap which is actually used in Iris
|
|
|
|
enum Color {
|
|
Red,
|
|
Black,
|
|
}
|
|
|
|
const { Red, Black } = Color;
|
|
|
|
type Comparator<K, V> = (a: [K, V], b: [K, V]) => number;
|
|
|
|
class TreeNode<K, V> {
|
|
key: K;
|
|
value: V;
|
|
left: TreeNode<K, V> | null = null;
|
|
right: TreeNode<K, V> | null = null;
|
|
color: Color = Color.Red; // Nodes are red by default when inserted
|
|
|
|
constructor(key: K, value: V) {
|
|
this.key = key;
|
|
this.value = value;
|
|
}
|
|
}
|
|
|
|
export default class RBSortedMap<K, V> {
|
|
private root: TreeNode<K, V> | null = null;
|
|
private compare: Comparator<K, V>;
|
|
|
|
constructor(initialEntries?: Iterable<readonly [K, V]>, compare?: string | Comparator<K, V>) {
|
|
if (compare) {
|
|
if (typeof compare === 'string') {
|
|
// comparison by object property
|
|
this.compare = (a, b) =>
|
|
a[1][compare] > b[1][compare] ? 1 : a[1][compare] < b[1][compare] ? -1 : 0;
|
|
} else {
|
|
this.compare = compare;
|
|
}
|
|
} else {
|
|
// comparison by key
|
|
this.compare = (a, b) => (a[0] > b[0] ? 1 : a[0] < b[0] ? -1 : 0);
|
|
}
|
|
|
|
if (initialEntries) {
|
|
for (const [key, value] of initialEntries) {
|
|
this.set(key, value);
|
|
}
|
|
}
|
|
}
|
|
|
|
private rotateLeft(node: TreeNode<K, V>): TreeNode<K, V> {
|
|
const newRoot = node.right!;
|
|
node.right = newRoot.left;
|
|
newRoot.left = node;
|
|
newRoot.color = node.color;
|
|
node.color = Red;
|
|
return newRoot;
|
|
}
|
|
|
|
private rotateRight(node: TreeNode<K, V>): TreeNode<K, V> {
|
|
const newRoot = node.left!;
|
|
node.left = newRoot.right;
|
|
newRoot.right = node;
|
|
newRoot.color = node.color;
|
|
node.color = Red;
|
|
return newRoot;
|
|
}
|
|
|
|
private flipColors(node: TreeNode<K, V>): void {
|
|
node.color = node.color === Black ? Red : Black;
|
|
if (node.left) node.left.color = node.left.color === Black ? Red : Black;
|
|
if (node.right) node.right.color = node.right.color === Black ? Red : Black;
|
|
}
|
|
|
|
private fixUp(node: TreeNode<K, V>): TreeNode<K, V> {
|
|
if (node.right && node.right.color === Red && (!node.left || node.left.color === Black)) {
|
|
node = this.rotateLeft(node);
|
|
}
|
|
|
|
if (node.left && node.left.color === Red && node.left.left && node.left.left.color === Red) {
|
|
node = this.rotateRight(node);
|
|
}
|
|
|
|
if (node.left && node.left.color === Red && node.right && node.right.color === Red) {
|
|
this.flipColors(node);
|
|
}
|
|
|
|
return node;
|
|
}
|
|
|
|
private insertNode(node: TreeNode<K, V> | null, key: K, value: V): TreeNode<K, V> {
|
|
if (node === null) {
|
|
return new TreeNode(key, value);
|
|
}
|
|
|
|
const comparison = this.compare([key, value], [node.key, node.value]);
|
|
|
|
if (comparison < 0) {
|
|
node.left = this.insertNode(node.left, key, value);
|
|
} else if (comparison > 0) {
|
|
node.right = this.insertNode(node.right, key, value);
|
|
} else {
|
|
node.value = value;
|
|
}
|
|
|
|
return this.fixUp(node);
|
|
}
|
|
|
|
set(key: K, value: V) {
|
|
this.root = this.insertNode(this.root, key, value);
|
|
this.root!.color = Black; // Root is always black
|
|
}
|
|
|
|
private findNode(node: TreeNode<K, V> | null, key: K): TreeNode<K, V> | null {
|
|
if (node === null) return null;
|
|
|
|
const comparison = this.compare([key, node.value], [node.key, node.value]);
|
|
|
|
if (comparison < 0) return this.findNode(node.left, key);
|
|
if (comparison > 0) return this.findNode(node.right, key);
|
|
|
|
return node;
|
|
}
|
|
|
|
get(key: K): V | undefined {
|
|
const node = this.findNode(this.root, key);
|
|
return node ? node.value : undefined;
|
|
}
|
|
|
|
first(): [K, V] | undefined {
|
|
let node = this.root;
|
|
while (node && node.left) {
|
|
node = node.left;
|
|
}
|
|
return node ? [node.key, node.value] : undefined;
|
|
}
|
|
|
|
last(): [K, V] | undefined {
|
|
let node = this.root;
|
|
while (node && node.right) {
|
|
node = node.right;
|
|
}
|
|
return node ? [node.key, node.value] : undefined;
|
|
}
|
|
|
|
private *inOrderTraversal(node: TreeNode<K, V> | null): IterableIterator<[K, V]> {
|
|
if (node) {
|
|
yield* this.inOrderTraversal(node.left);
|
|
yield [node.key, node.value];
|
|
yield* this.inOrderTraversal(node.right);
|
|
}
|
|
}
|
|
|
|
*[Symbol.iterator](): Iterator<[K, V]> {
|
|
yield* this.inOrderTraversal(this.root);
|
|
}
|
|
|
|
*reverse(): Iterator<[K, V]> {
|
|
// For simplicity, we use an array to reverse
|
|
const arr = [...this.inOrderTraversal(this.root)];
|
|
for (let i = arr.length - 1; i >= 0; i--) {
|
|
yield arr[i];
|
|
}
|
|
}
|
|
|
|
*keys(): IterableIterator<K> {
|
|
for (const entry of this.inOrderTraversal(this.root)) {
|
|
yield entry[0];
|
|
}
|
|
}
|
|
|
|
*values(): IterableIterator<V> {
|
|
for (const entry of this.inOrderTraversal(this.root)) {
|
|
yield entry[1];
|
|
}
|
|
}
|
|
|
|
*entries(): IterableIterator<[K, V]> {
|
|
for (const entry of this.inOrderTraversal(this.root)) {
|
|
yield entry;
|
|
}
|
|
}
|
|
|
|
range(
|
|
options: {
|
|
gte?: K;
|
|
lte?: K;
|
|
direction?: 'asc' | 'desc';
|
|
} = {},
|
|
): IterableIterator<[K, V]> {
|
|
const { gte, lte, direction = 'asc' } = options;
|
|
|
|
if (direction === 'asc') {
|
|
return this.inOrderRange(this.root, gte, lte);
|
|
} else {
|
|
return this.reverseInOrderRange(this.root, gte, lte);
|
|
}
|
|
}
|
|
|
|
private *inOrderRange(node: TreeNode<K, V> | null, gte?: K, lte?: K): Generator<[K, V]> {
|
|
if (!node) return;
|
|
|
|
if (gte) {
|
|
const gteComparison = this.compare([gte, this.get(gte) as V], [node.key, node.value]);
|
|
if (gteComparison > 0) {
|
|
yield* this.inOrderRange(node.right, gte, lte);
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (lte) {
|
|
const lteComparison = this.compare([lte, this.get(lte) as V], [node.key, node.value]);
|
|
if (lteComparison < 0) {
|
|
yield* this.inOrderRange(node.left, gte, lte);
|
|
return;
|
|
}
|
|
}
|
|
|
|
yield* this.inOrderRange(node.left, gte, lte);
|
|
yield [node.key, node.value];
|
|
yield* this.inOrderRange(node.right, gte, lte);
|
|
}
|
|
|
|
private *reverseInOrderRange(node: TreeNode<K, V> | null, gte?: K, lte?: K): Generator<[K, V]> {
|
|
if (!node) return;
|
|
|
|
if (lte) {
|
|
const lteComparison = this.compare([lte, this.get(lte) as V], [node.key, node.value]);
|
|
if (lteComparison < 0) {
|
|
yield* this.reverseInOrderRange(node.left, gte, lte);
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (gte) {
|
|
const gteComparison = this.compare([gte, this.get(gte) as V], [node.key, node.value]);
|
|
if (gteComparison > 0) {
|
|
yield* this.reverseInOrderRange(node.right, gte, lte);
|
|
return;
|
|
}
|
|
}
|
|
|
|
yield [node.key, node.value];
|
|
yield* this.reverseInOrderRange(node.right, gte, lte);
|
|
yield* this.reverseInOrderRange(node.left, gte, lte);
|
|
}
|
|
|
|
has(key: K): boolean {
|
|
return this.findNode(this.root, key) !== null;
|
|
}
|
|
|
|
private isRed(node: TreeNode<K, V> | null): boolean {
|
|
if (!node) return false;
|
|
return node.color === Red;
|
|
}
|
|
|
|
private moveRedLeft(node: TreeNode<K, V>): TreeNode<K, V> {
|
|
this.flipColors(node);
|
|
if (node.right && this.isRed(node.right.left)) {
|
|
node.right = this.rotateRight(node.right);
|
|
node = this.rotateLeft(node);
|
|
this.flipColors(node);
|
|
}
|
|
return node;
|
|
}
|
|
|
|
private moveRedRight(node: TreeNode<K, V>): TreeNode<K, V> {
|
|
this.flipColors(node);
|
|
if (node.left && this.isRed(node.left.left)) {
|
|
node = this.rotateRight(node);
|
|
this.flipColors(node);
|
|
}
|
|
return node;
|
|
}
|
|
|
|
private deleteMinNode(node: TreeNode<K, V>): TreeNode<K, V> | null {
|
|
if (!node.left) return null;
|
|
if (!this.isRed(node.left) && !this.isRed(node.left.left)) {
|
|
node = this.moveRedLeft(node);
|
|
}
|
|
node.left = this.deleteMinNode(node.left!);
|
|
return this.fixUp(node);
|
|
}
|
|
|
|
private minNode(node: TreeNode<K, V>): TreeNode<K, V> {
|
|
while (node.left) {
|
|
node = node.left;
|
|
}
|
|
return node;
|
|
}
|
|
|
|
private deleteNode(node: TreeNode<K, V> | null, key: K): TreeNode<K, V> | null {
|
|
if (this.compare([key, null!], [node!.key, node!.value]) < 0) {
|
|
if (!node!.left) return null;
|
|
if (!this.isRed(node!.left) && !this.isRed(node!.left.left)) {
|
|
node = this.moveRedLeft(node!);
|
|
}
|
|
node!.left = this.deleteNode(node!.left, key);
|
|
} else {
|
|
if (this.isRed(node!.left)) {
|
|
node = this.rotateRight(node!);
|
|
}
|
|
if (this.compare([key, null!], [node!.key, node!.value]) === 0 && !node!.right) {
|
|
return null;
|
|
}
|
|
if (!this.isRed(node!.right) && !this.isRed(node!.right!.left)) {
|
|
node = this.moveRedRight(node!);
|
|
}
|
|
if (this.compare([key, null!], [node!.key, node!.value]) === 0) {
|
|
const x = this.minNode(node!.right!);
|
|
node!.key = x.key;
|
|
node!.value = x.value;
|
|
node!.right = this.deleteMinNode(node!.right!);
|
|
} else {
|
|
node!.right = this.deleteNode(node!.right, key);
|
|
}
|
|
}
|
|
return this.fixUp(node!);
|
|
}
|
|
|
|
delete(key: K): boolean {
|
|
if (!this.has(key)) return false;
|
|
if (!this.isRed(this.root!.left) && !this.isRed(this.root!.right)) {
|
|
this.root!.color = Red;
|
|
}
|
|
this.root = this.deleteNode(this.root, key);
|
|
if (this.root) this.root!.color = Black;
|
|
return true;
|
|
}
|
|
|
|
// Calculating size efficiently would require augmenting the tree with size data.
|
|
get size(): number {
|
|
// Naive implementation
|
|
return [...this.keys()].length;
|
|
}
|
|
}
|