Created
June 26, 2018 01:37
-
-
Save hackerwins/f51fdcb4b30e6f356a98f8d4bfda0b95 to your computer and use it in GitHub Desktop.
SortedMap using LLRB Tree.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
/** | |
* Implementation of an SortedMap using Left-leaning Red-Black Tree | |
* | |
* Original paper on Left-leaning Red-Black Trees: | |
* - http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf | |
* | |
* Invariant 1: No red node has a red child | |
* Invariant 2: Every leaf path has the same number of black nodes | |
* Invariant 3: Only the left child can be red (left leaning) | |
*/ | |
interface Entry<K, V> { | |
key: K; | |
value: V; | |
} | |
class LLRBNode<K, V> { | |
public key: K; | |
public value: V; | |
public left: LLRBNode<K, V>; | |
public right: LLRBNode<K, V>; | |
public isRed: boolean; | |
constructor(key: K, value: V, isRed: boolean) { | |
this.key = key; | |
this.value = value; | |
this.isRed = isRed; | |
this.left = null; | |
this.right = null; | |
} | |
} | |
export type Comparator<K> = (keyA: K, keyB: K) => number; | |
const DEFAULT_COMPARATOR = (a, b) => { | |
if (a === b) { | |
return 0; | |
} else if (a < b) { | |
return -1; | |
} else { | |
return 1; | |
} | |
}; | |
export class SortedMapIterator<K, V> { | |
public stack: Array<Entry<K, V>>; | |
constructor(root: LLRBNode<K, V>) { | |
this.stack = []; | |
this.traverseInorder(root); | |
} | |
// TODO: Replace with iterative approach, if we encounter performance problem. | |
private traverseInorder(node: LLRBNode<K, V>): void { | |
if (node === null) { | |
return; | |
} | |
this.traverseInorder(node.left); | |
this.stack.push({ | |
key: node.key, | |
value: node.value | |
}); | |
this.traverseInorder(node.right); | |
} | |
} | |
export class SortedMap<K, V> { | |
private root: LLRBNode<K, V>; | |
private comparator: Comparator<K>; | |
constructor(comparator?: Comparator<K>) { | |
this.root = null; | |
this.comparator = typeof comparator !== 'undefined' ? comparator : DEFAULT_COMPARATOR; | |
} | |
public put(key: K, value: V): V { | |
this.root = this.putInternal(this.root, key, value); | |
this.root.isRed = false; | |
return value; | |
} | |
public get(key: K): V { | |
const node = this.getInternal(this.root, key); | |
return node !== null ? node.value : null; | |
} | |
public remove(key: K): void { | |
if (!this.isRed(this.root.left) && !this.isRed(this.root.right)) { | |
this.root.isRed = true; | |
} | |
this.root = this.removeInternal(this.root, key); | |
if (this.root !== null) { | |
this.root.isRed = false; | |
} | |
} | |
public getIterator(): SortedMapIterator<K, V> { | |
return new SortedMapIterator(this.root); | |
} | |
public values(): V[] { | |
const values = []; | |
for (const entry of this.getIterator().stack) { | |
values.push(entry.value); | |
} | |
return values; | |
} | |
private getInternal(node: LLRBNode<K, V>, key: K): LLRBNode<K, V> { | |
while (node !== null) { | |
const compare = this.comparator(key, node.key); | |
if (compare === 0) { | |
return node; | |
} else if (compare < 0) { | |
node = node.left; | |
} else if (compare > 0) { | |
node = node.right; | |
} | |
} | |
return null; | |
} | |
private putInternal(node: LLRBNode<K, V>, key: K, value: V): LLRBNode<K, V> { | |
if (node === null) { | |
return new LLRBNode(key, value, true); | |
} | |
const compare = this.comparator(key, node.key); | |
if (compare < 0) { | |
node.left = this.putInternal(node.left, key, value); | |
} else if (compare > 0) { | |
node.right = this.putInternal(node.right, key, value); | |
} else { | |
node.value = value; | |
} | |
if (this.isRed(node.right) && !this.isRed(node.left)) { | |
node = this.rotateLeft(node); | |
} | |
if (this.isRed(node.left) && this.isRed(node.left.left)) { | |
node = this.rotateRight(node); | |
} | |
if (this.isRed(node.left) && this.isRed(node.right)) { | |
this.flipColors(node); | |
} | |
return node; | |
} | |
private removeInternal(node: LLRBNode<K, V>, key: K): LLRBNode<K, V> { | |
if (this.comparator(key, node.key) < 0) { | |
if (!this.isRed(node.left) && !this.isRed(node.left.left)) { | |
node = this.moveRedLeft(node); | |
} | |
node.left = this.removeInternal(node.left, key); | |
} else { | |
if (this.isRed(node.left)) { | |
node = this.rotateRight(node); | |
} | |
if (this.comparator(key, node.key) === 0 && node.right === null) { | |
return null; | |
} | |
if (!this.isRed(node.right) && !this.isRed(node.right.left)) { | |
node = this.moveRedRight(node); | |
} | |
if (this.comparator(key, node.key) === 0) { | |
const smallest = this.min(node.right); | |
node.value = smallest.value; | |
node.key = smallest.key; | |
node.right = this.removeMin(node.right); | |
} else { | |
node.right = this.removeInternal(node.right, key); | |
} | |
} | |
return this.fixUp(node); | |
} | |
private min(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
if (node.left == null) { | |
return node; | |
} else { | |
return this.min(node.left); | |
} | |
} | |
private removeMin(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
if (node.left === null) { | |
return null; | |
} | |
if (!this.isRed(node.left) && !this.isRed(node.left.left)) { | |
node = this.moveRedLeft(node); | |
} | |
node.left = this.removeMin(node.left); | |
return this.fixUp(node); | |
} | |
private fixUp(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
if (this.isRed(node.right)) { | |
node = this.rotateLeft(node); | |
} | |
if (this.isRed(node.left) && this.isRed(node.left.left)) { | |
node = this.rotateRight(node); | |
} | |
if (this.isRed(node.left) && this.isRed(node.right)) { | |
this.flipColors(node); | |
} | |
return node; | |
} | |
private moveRedLeft(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
this.flipColors(node); | |
if (this.isRed(node.right.left)) { | |
node.right = this.rotateRight(node.right); | |
node = this.rotateLeft(node); | |
this.flipColors(node); | |
} | |
return node; | |
} | |
private moveRedRight(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
this.flipColors(node); | |
if (this.isRed(node.left.left)) { | |
node = this.rotateRight(node); | |
this.flipColors(node); | |
} | |
return node; | |
} | |
private isRed(node: LLRBNode<K, V>): boolean { | |
return node !== null && node.isRed; | |
} | |
private rotateLeft(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
const x = node.right; | |
node.right = x.left; | |
x.left = node; | |
x.isRed = x.left.isRed; | |
x.left.isRed = true; | |
return x; | |
} | |
private rotateRight(node: LLRBNode<K, V>): LLRBNode<K, V> { | |
const x = node.left; | |
node.left = x.right; | |
x.right = node; | |
x.isRed = x.right.isRed; | |
x.right.isRed = true; | |
return x; | |
} | |
private flipColors(node: LLRBNode<K, V>): void { | |
node.isRed = !node.isRed; | |
node.left.isRed = !node.left.isRed; | |
node.right.isRed = !node.right.isRed; | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment