685 lines
15 KiB
TypeScript
685 lines
15 KiB
TypeScript
import Node from "./node";
|
|
import { Key, Value } from "./types";
|
|
|
|
export type Comparator<Key> = (a: Key, b: Key) => number;
|
|
export type Visitor<Key, Value> = (
|
|
node: Node<Key, Value>
|
|
) => undefined | boolean | void;
|
|
export type NodePrinter<Key, Value> = (node: Node<Key, Value>) => string;
|
|
|
|
/* follows "An implementation of top-down splaying"
|
|
* by D. Sleator <sleator@cs.cmu.edu> March 1992
|
|
*/
|
|
|
|
function DEFAULT_COMPARE(a: Key, b: Key): number {
|
|
return a > b ? 1 : a < b ? -1 : 0;
|
|
}
|
|
|
|
type TreeNodeList<Key, Value> = { head: Node<Key, Value> | null };
|
|
|
|
/**
|
|
* Simple top down splay, not requiring i to be in the tree t.
|
|
*/
|
|
function splay(
|
|
i: Key,
|
|
t: Node<Key, Value>,
|
|
comparator: Comparator<Key>
|
|
): Node<Key, Value> {
|
|
const N = new Node(null, null);
|
|
let l = N;
|
|
let r = N;
|
|
|
|
while (true) {
|
|
const cmp = comparator(i, t!.key);
|
|
//if (i < t.key) {
|
|
if (cmp < 0) {
|
|
if (t!.left === null) break;
|
|
//if (i < t.left.key) {
|
|
if (comparator(i, t.left.key) < 0) {
|
|
const y = t.left; /* rotate right */
|
|
t.left = y.right;
|
|
y.right = t;
|
|
t = y;
|
|
if (t.left === null) break;
|
|
}
|
|
r.left = t; /* link right */
|
|
r = t;
|
|
t = t.left;
|
|
//} else if (i > t.key) {
|
|
} else if (cmp > 0) {
|
|
if (t.right === null) break;
|
|
//if (i > t.right.key) {
|
|
if (comparator(i, t.right.key) > 0) {
|
|
const y = t.right; /* rotate left */
|
|
t.right = y.left;
|
|
y.left = t;
|
|
t = y;
|
|
if (t.right === null) break;
|
|
}
|
|
l.right = t; /* link left */
|
|
l = t;
|
|
t = t.right;
|
|
} else break;
|
|
}
|
|
/* assemble */
|
|
l.right = t.left;
|
|
r.left = t.right;
|
|
t.left = N.right;
|
|
t.right = N.left;
|
|
return t;
|
|
}
|
|
|
|
function insert(
|
|
i: Key,
|
|
data: Value,
|
|
t: Node<Key, Value>,
|
|
comparator: Comparator<Key>
|
|
): Node<Key, Value> {
|
|
const node = new Node(i, data);
|
|
|
|
if (t === null) {
|
|
node.left = node.right = null;
|
|
return node;
|
|
}
|
|
|
|
t = splay(i, t, comparator);
|
|
const cmp = comparator(i, t.key);
|
|
if (cmp < 0) {
|
|
node.left = t.left;
|
|
node.right = t;
|
|
t.left = null;
|
|
} else if (cmp >= 0) {
|
|
node.right = t.right;
|
|
node.left = t;
|
|
t.right = null;
|
|
}
|
|
return node;
|
|
}
|
|
|
|
function split(
|
|
key: Key,
|
|
v: Node<Key, Value>,
|
|
comparator: Comparator<Key>
|
|
): {
|
|
left: Node<Key, Value> | null;
|
|
right: Node<Key, Value> | null;
|
|
} {
|
|
let left = null;
|
|
let right = null;
|
|
if (v) {
|
|
v = splay(key, v, comparator);
|
|
|
|
const cmp = comparator(v.key, key);
|
|
if (cmp === 0) {
|
|
left = v.left;
|
|
right = v.right;
|
|
} else if (cmp < 0) {
|
|
right = v.right;
|
|
v.right = null;
|
|
left = v;
|
|
} else {
|
|
left = v.left;
|
|
v.left = null;
|
|
right = v;
|
|
}
|
|
}
|
|
return { left, right };
|
|
}
|
|
|
|
function merge(
|
|
left: Node<Key, Value> | null,
|
|
right: Node<Key, Value> | null,
|
|
comparator: Comparator<Key>
|
|
) {
|
|
if (right === null) return left;
|
|
if (left === null) return right;
|
|
|
|
right = splay(left.key, right, comparator);
|
|
right.left = left;
|
|
return right;
|
|
}
|
|
|
|
type StringCollector = (s: string) => void;
|
|
|
|
/**
|
|
* Prints level of the tree
|
|
*/
|
|
function printRow(
|
|
root: Node<Key, Value>,
|
|
prefix: string,
|
|
isTail: boolean,
|
|
out: StringCollector,
|
|
printNode: NodePrinter<Key, Value>
|
|
) {
|
|
if (root) {
|
|
out(`${prefix}${isTail ? "└── " : "├── "}${printNode(root)}\n`);
|
|
const indent = prefix + (isTail ? " " : "│ ");
|
|
if (root.left) printRow(root.left, indent, false, out, printNode);
|
|
if (root.right) printRow(root.right, indent, true, out, printNode);
|
|
}
|
|
}
|
|
|
|
export default class Tree<Key = number, Value = any> {
|
|
private _comparator: Comparator<Key>;
|
|
private _root: Node<Key, Value> | null = null;
|
|
private _size: number = 0;
|
|
|
|
constructor(comparator = DEFAULT_COMPARE) {
|
|
this._comparator = comparator;
|
|
}
|
|
|
|
/**
|
|
* Inserts a key, allows duplicates
|
|
*/
|
|
public insert(key: Key, data?: Value): Node<Key, Value> {
|
|
this._size++;
|
|
return (this._root = insert(key, data, this._root!, this._comparator));
|
|
}
|
|
|
|
/**
|
|
* Adds a key, if it is not present in the tree
|
|
*/
|
|
public add(key: Key, data?: Value): Node<Key, Value> {
|
|
const node = new Node(key, data);
|
|
|
|
if (this._root === null) {
|
|
node.left = node.right = null;
|
|
this._size++;
|
|
this._root = node;
|
|
}
|
|
|
|
const comparator = this._comparator;
|
|
const t = splay(key, this._root, comparator);
|
|
const cmp = comparator(key, t.key);
|
|
if (cmp === 0) this._root = t;
|
|
else {
|
|
if (cmp < 0) {
|
|
node.left = t.left;
|
|
node.right = t;
|
|
t.left = null;
|
|
} else if (cmp > 0) {
|
|
node.right = t.right;
|
|
node.left = t;
|
|
t.right = null;
|
|
}
|
|
this._size++;
|
|
this._root = node;
|
|
}
|
|
|
|
return this._root;
|
|
}
|
|
|
|
/**
|
|
* @param {Key} key
|
|
* @return {Node|null}
|
|
*/
|
|
public remove(key: Key): void {
|
|
this._root = this._remove(key, this._root!, this._comparator);
|
|
}
|
|
|
|
/**
|
|
* Deletes i from the tree if it's there
|
|
*/
|
|
private _remove(i: Key, t: Node<Key, Value>, comparator: Comparator<Key>) {
|
|
let x;
|
|
if (t === null) return null;
|
|
t = splay(i, t, comparator);
|
|
const cmp = comparator(i, t.key);
|
|
if (cmp === 0) {
|
|
/* found it */
|
|
if (t.left === null) {
|
|
x = t.right!;
|
|
} else {
|
|
x = splay(i, t.left, comparator);
|
|
x.right = t.right;
|
|
}
|
|
this._size--;
|
|
return x;
|
|
}
|
|
return t; /* It wasn't there */
|
|
}
|
|
|
|
/**
|
|
* Removes and returns the node with smallest key
|
|
*/
|
|
public pop(): { key: Key; data: Value } | null {
|
|
let node = this._root;
|
|
if (node) {
|
|
while (node.left) node = node.left;
|
|
this._root = splay(node.key, this._root!, this._comparator);
|
|
this._root = this._remove(node.key, this._root, this._comparator);
|
|
return { key: node.key, data: node.data };
|
|
}
|
|
return null;
|
|
}
|
|
|
|
/**
|
|
* Find without splaying
|
|
*/
|
|
public findStatic(key: Key): Node<Key, Value> | null {
|
|
let current = this._root;
|
|
const compare = this._comparator;
|
|
while (current) {
|
|
const cmp = compare(key, current.key);
|
|
if (cmp === 0) return current;
|
|
else if (cmp < 0) current = current.left;
|
|
else current = current.right;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
public find(key: Key): Node<Key, Value> | null {
|
|
if (this._root) {
|
|
this._root = splay(key, this._root, this._comparator);
|
|
if (this._comparator(key, this._root.key) !== 0) return null;
|
|
}
|
|
return this._root;
|
|
}
|
|
|
|
public contains(key: Key): boolean {
|
|
let current = this._root;
|
|
const compare = this._comparator;
|
|
while (current) {
|
|
const cmp = compare(key, current.key);
|
|
if (cmp === 0) return true;
|
|
else if (cmp < 0) current = current.left;
|
|
else current = current.right;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
public forEach(visitor: Visitor<Key, Value>, ctx?: any): Tree<Key, Value> {
|
|
let current = this._root;
|
|
const Q = []; /* Initialize stack s */
|
|
let done = false;
|
|
|
|
while (!done) {
|
|
if (current !== null) {
|
|
Q.push(current);
|
|
current = current.left;
|
|
} else {
|
|
if (Q.length !== 0) {
|
|
current = Q.pop()!;
|
|
visitor.call(ctx, current);
|
|
|
|
current = current.right;
|
|
} else done = true;
|
|
}
|
|
}
|
|
return this;
|
|
}
|
|
|
|
/**
|
|
* Walk key range from `low` to `high`. Stops if `fn` returns a value.
|
|
*/
|
|
public range(
|
|
low: Key,
|
|
high: Key,
|
|
fn: Visitor<Key, Value>,
|
|
ctx?: any
|
|
): Tree<Key, Value> {
|
|
const Q = [];
|
|
const compare = this._comparator;
|
|
let node = this._root;
|
|
let cmp;
|
|
|
|
while (Q.length !== 0 || node) {
|
|
if (node) {
|
|
Q.push(node);
|
|
node = node.left;
|
|
} else {
|
|
node = Q.pop()!;
|
|
cmp = compare(node.key, high);
|
|
if (cmp > 0) {
|
|
break;
|
|
} else if (compare(node.key, low) >= 0) {
|
|
if (fn.call(ctx, node)) return this; // stop if smth is returned
|
|
}
|
|
node = node.right;
|
|
}
|
|
}
|
|
return this;
|
|
}
|
|
|
|
/**
|
|
* Returns array of keys
|
|
*/
|
|
public keys(): Key[] {
|
|
const keys: Key[] = [];
|
|
this.forEach(({ key }) => {
|
|
keys.push(key);
|
|
});
|
|
return keys;
|
|
}
|
|
|
|
/**
|
|
* Returns array of all the data in the nodes
|
|
*/
|
|
public values(): Value[] {
|
|
const values: Value[] = [];
|
|
this.forEach(({ data }) => {
|
|
values.push(data);
|
|
});
|
|
return values;
|
|
}
|
|
|
|
public min(): Key | null {
|
|
if (this._root) return this.minNode(this._root)!.key;
|
|
return null;
|
|
}
|
|
|
|
public max(): Key | null {
|
|
if (this._root) return this.maxNode(this._root)!.key;
|
|
return null;
|
|
}
|
|
|
|
public minNode(t = this._root): Node<Key, Value> | null {
|
|
if (t) while (t.left) t = t.left;
|
|
return t;
|
|
}
|
|
|
|
public maxNode(t = this._root): Node<Key, Value> | null {
|
|
if (t) while (t.right) t = t.right;
|
|
return t;
|
|
}
|
|
|
|
/**
|
|
* Returns node at given index
|
|
*/
|
|
public at(index: number): Node<Key, Value> | null {
|
|
let current = this._root;
|
|
let done = false;
|
|
let i = 0;
|
|
const Q = [];
|
|
|
|
while (!done) {
|
|
if (current) {
|
|
Q.push(current);
|
|
current = current.left;
|
|
} else {
|
|
if (Q.length > 0) {
|
|
current = Q.pop()!;
|
|
if (i === index) return current;
|
|
i++;
|
|
current = current.right!;
|
|
} else done = true;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
public next(d: Node<Key, Value>): Node<Key, Value> | null {
|
|
let root = this._root;
|
|
let successor = null;
|
|
|
|
if (d.right) {
|
|
successor = d.right;
|
|
while (successor.left) successor = successor.left;
|
|
return successor;
|
|
}
|
|
|
|
const comparator = this._comparator;
|
|
while (root) {
|
|
const cmp = comparator(d.key, root.key);
|
|
if (cmp === 0) break;
|
|
else if (cmp < 0) {
|
|
successor = root;
|
|
root = root.left;
|
|
} else root = root.right;
|
|
}
|
|
|
|
return successor;
|
|
}
|
|
|
|
public prev(d: Node<Key, Value>): Node<Key, Value> | null {
|
|
let root = this._root;
|
|
let predecessor = null;
|
|
|
|
if (d.left !== null) {
|
|
predecessor = d.left;
|
|
while (predecessor.right) predecessor = predecessor.right;
|
|
return predecessor;
|
|
}
|
|
|
|
const comparator = this._comparator;
|
|
while (root) {
|
|
const cmp = comparator(d.key, root.key);
|
|
if (cmp === 0) break;
|
|
else if (cmp < 0) root = root.left;
|
|
else {
|
|
predecessor = root;
|
|
root = root.right;
|
|
}
|
|
}
|
|
return predecessor;
|
|
}
|
|
|
|
public clear(): Tree<Key, Value> {
|
|
this._root = null;
|
|
this._size = 0;
|
|
return this;
|
|
}
|
|
|
|
public toList() {
|
|
return toList(this._root!);
|
|
}
|
|
|
|
/**
|
|
* Bulk-load items. Both array have to be same size
|
|
*/
|
|
public load(keys: Key[], values: Value[] = [], presort: boolean = false) {
|
|
let size = keys.length;
|
|
const comparator = this._comparator;
|
|
|
|
// sort if needed
|
|
if (presort) sort(keys, values, 0, size - 1, comparator);
|
|
|
|
if (this._root === null) {
|
|
// empty tree
|
|
this._root = loadRecursive(keys, values, 0, size);
|
|
this._size = size;
|
|
} else {
|
|
// that re-builds the whole tree from two in-order traversals
|
|
const mergedList = mergeLists(
|
|
this.toList(),
|
|
createList(keys, values),
|
|
comparator
|
|
);
|
|
size = this._size + size;
|
|
this._root = sortedListToBST({ head: mergedList }, 0, size);
|
|
}
|
|
return this;
|
|
}
|
|
|
|
public isEmpty(): boolean {
|
|
return this._root === null;
|
|
}
|
|
|
|
get size(): number {
|
|
return this._size;
|
|
}
|
|
get root(): Node<Key, Value> | null {
|
|
return this._root;
|
|
}
|
|
|
|
public toString(
|
|
printNode: NodePrinter<Key, Value> = (n) => String(n.key)
|
|
): string {
|
|
const out: string[] = [];
|
|
printRow(this._root!, "", true, (v) => out.push(v), printNode);
|
|
return out.join("");
|
|
}
|
|
|
|
public update(key: Key, newKey: Key, newData?: Value): void {
|
|
const comparator = this._comparator;
|
|
let { left, right } = split(key, this._root!, comparator);
|
|
if (comparator(key, newKey) < 0) {
|
|
right = insert(newKey, newData, right!, comparator);
|
|
} else {
|
|
left = insert(newKey, newData, left!, comparator);
|
|
}
|
|
this._root = merge(left, right, comparator);
|
|
}
|
|
|
|
public split(key: Key) {
|
|
return split(key, this._root!, this._comparator);
|
|
}
|
|
|
|
*[Symbol.iterator]() {
|
|
let current = this._root;
|
|
const Q: Node<Key, Value>[] = []; /* Initialize stack s */
|
|
let done = false;
|
|
|
|
while (!done) {
|
|
if (current !== null) {
|
|
Q.push(current);
|
|
current = current.left;
|
|
} else {
|
|
if (Q.length !== 0) {
|
|
current = Q.pop()!;
|
|
yield current;
|
|
|
|
current = current.right;
|
|
} else done = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
function loadRecursive(
|
|
keys: Key[],
|
|
values: Value[],
|
|
start: number,
|
|
end: number
|
|
): Node<Key, Value> | null {
|
|
const size = end - start;
|
|
if (size > 0) {
|
|
const middle = start + Math.floor(size / 2);
|
|
const key = keys[middle];
|
|
const data = values[middle];
|
|
const node = new Node(key, data);
|
|
node.left = loadRecursive(keys, values, start, middle);
|
|
node.right = loadRecursive(keys, values, middle + 1, end);
|
|
return node;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
function createList(keys: Key[], values: Value[]): Node<Key, Value> {
|
|
const head = new Node<Key, Value>(null, null);
|
|
let p: Node<Key, Value> = head;
|
|
for (let i = 0; i < keys.length; i++) {
|
|
p = p.next = new Node(keys[i], values[i]);
|
|
}
|
|
p.next = null;
|
|
return head.next!;
|
|
}
|
|
|
|
function toList(root: Node<Key, Value>): Node<Key, Value> {
|
|
let current = root;
|
|
const Q = [];
|
|
let done = false;
|
|
|
|
const head = new Node<Key, Value>(null, null);
|
|
let p = head;
|
|
|
|
while (!done) {
|
|
if (current) {
|
|
Q.push(current);
|
|
current = current.left!;
|
|
} else {
|
|
if (Q.length > 0) {
|
|
current = p = p.next = Q.pop()!;
|
|
current = current.right!;
|
|
} else done = true;
|
|
}
|
|
}
|
|
p.next = null; // that'll work even if the tree was empty
|
|
return head.next!;
|
|
}
|
|
|
|
function sortedListToBST(
|
|
list: TreeNodeList<Key, Value>,
|
|
start: number,
|
|
end: number
|
|
) {
|
|
const size = end - start;
|
|
if (size > 0) {
|
|
const middle = start + Math.floor(size / 2);
|
|
const left = sortedListToBST(list, start, middle);
|
|
|
|
const root = list.head;
|
|
root!.left = left;
|
|
|
|
list.head = list.head!.next;
|
|
|
|
root!.right = sortedListToBST(list, middle + 1, end);
|
|
return root!;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
function mergeLists<Key, Value>(
|
|
l1: Node<Key, Value>,
|
|
l2: Node<Key, Value>,
|
|
compare: Comparator<Key>
|
|
): Node<Key, Value> {
|
|
const head: Node<Key, Value> = new Node<Key, Value>(null as Key, null); // dummy
|
|
let p = head;
|
|
|
|
let p1: Node<Key, Value> = l1;
|
|
let p2: Node<Key, Value> = l2;
|
|
|
|
while (p1 !== null && p2 !== null) {
|
|
if (compare(p1.key, p2.key) < 0) {
|
|
p.next = p1;
|
|
p1 = p1.next!;
|
|
} else {
|
|
p.next = p2;
|
|
p2 = p2.next!;
|
|
}
|
|
p = p.next;
|
|
}
|
|
|
|
if (p1 !== null) {
|
|
p.next = p1;
|
|
} else if (p2 !== null) {
|
|
p.next = p2;
|
|
}
|
|
|
|
return head.next!;
|
|
}
|
|
|
|
function sort(
|
|
keys: Key[],
|
|
values: Value[],
|
|
left: number,
|
|
right: number,
|
|
compare: Comparator<Key>
|
|
) {
|
|
if (left >= right) return;
|
|
|
|
const pivot = keys[(left + right) >> 1];
|
|
let i = left - 1;
|
|
let j = right + 1;
|
|
|
|
while (true) {
|
|
do i++;
|
|
while (compare(keys[i], pivot) < 0);
|
|
do j--;
|
|
while (compare(keys[j], pivot) > 0);
|
|
if (i >= j) break;
|
|
|
|
let tmp = keys[i];
|
|
keys[i] = keys[j];
|
|
keys[j] = tmp;
|
|
|
|
tmp = values[i];
|
|
values[i] = values[j];
|
|
values[j] = tmp;
|
|
}
|
|
|
|
sort(keys, values, left, j, compare);
|
|
sort(keys, values, j + 1, right, compare);
|
|
}
|