RedBlackTree.ts

RedBlackTree.ts

type RedBlackTree<T> = RedBlackTreeNode<T> | RedBlackTreeLeaf;

interface RedBlackTreeNode<T> {
    left: RedBlackTree<T>;
    right: RedBlackTree<T>;
    color: Color;
    value: T;
}

type RedBlackTreeLeaf = null;

type Comparer<T> = (first: T, second: T) => number;

enum Color {
    RED,
    BLACK
};

const Nil: RedBlackTree<any> = null;

function insert<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): RedBlackTreeNode<T> {
    const node = insertStep(tree, value, comparer);
    return changeColor(node, Color.BLACK);
}

function insertStep<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): RedBlackTreeNode<T> {
    if (isLeaf(tree)) {
        return {
            left: Nil,
            right: Nil,
            color: Color.RED,
            value
        };
    } else {
        const ordering = comparer(value, tree.value);
        if (ordering < 0) {
            const left = insertStep(tree.left, value, comparer);
            if (tree.color === Color.BLACK) {
                return balance(left, tree.right, tree.value);
            } else {
                return { left, right: tree.right, color: Color.RED, value: tree.value };
            }
        } else if (ordering > 0) {
            const right = insertStep(tree.right, value, comparer);
            if (tree.color === Color.BLACK)  {
                return balance(tree.left, right, tree.value);
            } else {
                return { left: tree.left, right, color: Color.RED, value: tree.value };
            }
        } else {
            return {
                left: tree.left,
                right: tree.right,
                color: tree.color,
                value
            };
        }
    }
}

function deleteFrom<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): RedBlackTree<T> {
    const node = deleteStep(tree, value, comparer);
    if (isLeaf(node)) {
        return node;
    } else {
        return changeColor(node, Color.BLACK);
    }
}

function deleteStep<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): RedBlackTree<T> {
    if (isLeaf(tree)) {
        return Nil;
    } else {
        const ordering = comparer(value, tree.value);
        if (ordering < 0) {
            const left = deleteStep(tree.left, value, comparer);
            if (isBlackNode(tree.left)) {
                return deleteFixupLeft(left, tree.right, tree.value);
            } else {
                return { left, right: tree.right, color: Color.RED, value: tree.value };
            }
        } else if (ordering > 0) {
            const right = deleteStep(tree.right, value, comparer);
            if (isBlackNode(tree.right)) {
                return deleteFixupRight(tree.left, right, tree.value);
            } else {
                return { left: tree.left, right, color: Color.RED, value: tree.value };
            }
        } else {
            return append(tree.left, tree.right);
        }
    }
}

function deleteFixupLeft<T>(left: RedBlackTree<T>, right: RedBlackTree<T>, value: T): RedBlackTree<T> {
    if (isRedNode(left)) {
        return {
            left: changeColor(left, Color.BLACK),
            right,
            color: Color.RED,
            value
        };
    } else {
        if (isNode(right)) {
            if (right.color === Color.BLACK) {
                return balance(left, changeColor(right, Color.RED), value);
            } else {  // if (right.color === Color.RED)
                if (isBlackNode(right.left)) {
                    return {
                        left: {
                            left,
                            right: right.left.left,
                            color: Color.BLACK,
                            value
                        },
                        right: balance(right.left.right, tryChangeColor(right.right, Color.RED), right.value),
                        color: Color.RED,
                        value: right.left.value
                    };
                }
            }
        }
    }
    return { left, right, color: Color.RED, value };
}

function deleteFixupRight<T>(left: RedBlackTree<T>, right: RedBlackTree<T>, value: T): RedBlackTree<T> {
    if (isRedNode(right)) {
        return {
            left,
            right: changeColor(right, Color.BLACK),
            color: Color.RED,
            value
        };
    } else if (isNode(left)) {
        if (left.color === Color.BLACK) {
            return balance(changeColor(left, Color.RED), right, value);
        } else {  // if (left.color === Color.RED)
            if (isBlackNode(left.right)) {
                return {
                    left: balance(tryChangeColor(left.left, Color.RED), left.right.left, left.value),
                    right: {
                        left: left.right.right,
                        right,
                        color: Color.BLACK,
                        value
                    },
                    color: Color.RED,
                    value: left.right.value
                };
            }
        }
    }
    return { left, right, color: Color.RED, value };
}

function member<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): boolean {
    if (isLeaf(tree)) {
        return false;
    } else {
        const ordering = comparer(value, tree.value);
        if (ordering < 0) {
            return member(tree.left, value, comparer);
        } else if (ordering > 0) {
            return member(tree.right, value, comparer);
        } else {
            return true;
        }
    }
}

function min<T>(tree: RedBlackTree<T>, comparer: Comparer<T>): T | null {
    let value = null;
    while (isNode(tree)) {
        value = tree.value;
        tree = tree.left;
    }
    return value;
}

function max<T>(tree: RedBlackTree<T>, comparer: Comparer<T>): T | null {
    let value = null;
    while (isNode(tree)) {
        value = tree.value;
        tree = tree.right;
    }
    return value;
}

function fromIterable<T>(iterable: Iterable<T>, comparer: Comparer<T>): RedBlackTree<T> {
    let tree = Nil;
    for (const element of iterable) {
        tree = insert(tree, element, comparer);
    }
    return tree;
}

function toArray<T>(tree: RedBlackTree<T>): T[] {
    return toArrayStep(tree, []);
}

function toArrayStep<T>(tree: RedBlackTree<T>, output: T[]): T[] {
    if (isNode(tree)) {
        if (isNode(tree.left)) {
            toArrayStep(tree.left, output);
        }
        output.push(tree.value);
        if (isNode(tree.right)) {
            toArrayStep(tree.right, output);
        }
    }
    return output;
}

function* iterate<T>(tree: RedBlackTree<T>): Iterable<T> {
    if (isNode(tree)) {
        if (isNode(tree.left)) {
            yield* iterate(tree.left);
        }
        yield tree.value;
        if (isNode(tree.right)) {
            yield* iterate(tree.right);
        }
    }
}

function* iterateGreaterThan<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): Iterable<T> {
    if (isNode(tree)) {
        const ordering = comparer(tree.value, value);
        if (ordering < 0) {
            if (isNode(tree.right)) {
                yield* iterateGreaterThan(tree.right, value, comparer);
            }
        } else if (ordering > 0) {
            if (isNode(tree.left)) {
                yield* iterateGreaterThan(tree.left, value, comparer);
            }
            yield tree.value;
            if (isNode(tree.right)) {
                yield* iterateGreaterThan(tree.right, value, comparer);
            }
        } else {
            if (isNode(tree.right)) {
                yield* iterateGreaterThan(tree.right, value, comparer);
            }
        }
    }
}

function* iterateGreaterThanOrEqual<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): Iterable<T> {
    if (isNode(tree)) {
        const ordering = comparer(tree.value, value);
        if (ordering < 0) {
            if (isNode(tree.right)) {
                yield* iterateGreaterThanOrEqual(tree.right, value, comparer);
            }
        } else if (ordering > 0) {
            if (isNode(tree.left)) {
                yield* iterateGreaterThanOrEqual(tree.left, value, comparer);
            }
            yield tree.value;
            if (isNode(tree.right)) {
                yield* iterateGreaterThanOrEqual(tree.right, value, comparer);
            }
        } else {
            yield tree.value;
            if (isNode(tree.right)) {
                yield* iterateGreaterThanOrEqual(tree.right, value, comparer);
            }
        }
    }
}

function* iterateLessThan<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): Iterable<T> {
    if (isNode(tree)) {
        const ordering = comparer(tree.value, value);
        if (ordering < 0) {
            if (isNode(tree.left)) {
                yield* iterateLessThan(tree.left, value, comparer);
            }
            yield tree.value;
            if (isNode(tree.right)) {
                yield* iterateLessThan(tree.right, value, comparer);
            }
        } else if (ordering > 0) {
            if (isNode(tree.left)) {
                yield* iterateLessThan(tree.left, value, comparer);
            }
        } else {
            if (isNode(tree.left)) {
                yield* iterateLessThan(tree.left, value, comparer);
            }
        }
    }
}

function* iterateLessThanOrEqual<T>(tree: RedBlackTree<T>, value: T, comparer: Comparer<T>): Iterable<T> {
    if (isNode(tree)) {
        const ordering = comparer(tree.value, value);
        if (ordering < 0) {
            if (isNode(tree.left)) {
                yield* iterateLessThanOrEqual(tree.left, value, comparer);
            }
            yield tree.value;
            if (isNode(tree.right)) {
                yield* iterateLessThanOrEqual(tree.right, value, comparer);
            }
        } else if (ordering > 0) {
            if (isNode(tree.left)) {
                yield* iterateLessThanOrEqual(tree.left, value, comparer);
            }
        } else {
            yield tree.value;
            if (isNode(tree.left)) {
                yield* iterateLessThanOrEqual(tree.left, value, comparer);
            }
        }
    }
}

function append<T>(left: RedBlackTree<T>, right: RedBlackTree<T>): RedBlackTree<T> {
    if (isLeaf(left)) {
        return right;
    }

    if (isLeaf(right)) {
        return left;
    }

    if (left.color === Color.RED && right.color === Color.RED) {
        const tree = append(left.right, right.left);
        if (isRedNode(tree)) {
            return {
                left: {
                    left: left.left,
                    right: tree.left,
                    color: Color.RED,
                    value: left.value
                },
                right: {
                    left: tree.right,
                    right: right.right,
                    color: Color.RED,
                    value: right.value
                },
                color: Color.RED,
                value: tree.value
            };
        } else {
            return {
                left: left.left,
                right: {
                    left: tree,
                    right: right.right,
                    color: Color.RED,
                    value: right.value
                },
                color: Color.RED,
                value: left.value
            }
        }
    } else if (left.color === Color.BLACK && right.color === Color.BLACK) {
        const tree = append(left.right, right.left);
        if (isRedNode(tree)) {
            return {
                left: {
                    left: left.left,
                    right: tree.left,
                    color: Color.BLACK,
                    value: left.value
                },
                right: {
                    left: tree.right,
                    right: right.right,
                    color: Color.BLACK,
                    value: right.value
                },
                color: Color.RED,
                value: tree.value
            };
        } else {
            return deleteFixupLeft(
                left.left,
                { left: tree, right: right.right, color: Color.BLACK, value: right.value },
                left.value
            );
        }
    } else if (left.color === Color.RED) {
        return {
            left: left.left,
            right: append(left.right, right),
            color: Color.RED,
            value: left.value
        };
    } else {  // if (right.color === Color.RED)
        return {
            left: append(left, right.left),
            right: right.right,
            color: Color.RED,
            value: right.value
        };
    }
}

function balance<T>(left: RedBlackTree<T>, right: RedBlackTree<T>, value: T): RedBlackTreeNode<T> {
    if (isRedNode(left)) {
        if (isRedNode(right)) {
            return {
                left: changeColor(left, Color.BLACK),
                right: changeColor(right, Color.BLACK),
                color: Color.RED,
                value: value
            };
        } else if (isRedNode(left.left)) {
            return {
                left: { left: left.left.left, right: left.left.right, color: Color.BLACK, value: left.left.value },
                right: { left: left.right, right, color: Color.BLACK, value },
                color: Color.RED,
                value: left.value
            };
        } else if (isRedNode(left.right)) {
            return {
                left: { left: left.left, right: left.right.left, color: Color.BLACK, value: left.value },
                right: { left: left.right.right, right, color: Color.BLACK, value },
                color: Color.RED,
                value: left.right.value
            };
        }
    } else if (isRedNode(right)) {
        if (isRedNode(right.left)) {
            return {
                left: { left, right: right.left.left, color: Color.BLACK, value },
                right: { left: right.left.right, right: right.right, color: Color.BLACK, value: right.value },
                color: Color.RED,
                value: right.left.value
            };
        } else if (isRedNode(right.right)) {
            return {
                left: { left, right: right.left, color: Color.BLACK, value },
                right: { left: right.right.left, right: right.right.right, color: Color.BLACK, value: right.right.value },
                color: Color.RED,
                value: right.value
            };
        }
    }

    return { left, right, color: Color.BLACK, value };
}

function changeColor(node: RedBlackTreeNode<any>, color: Color): RedBlackTreeNode<any> {
    return {
        left: node.left,
        right: node.right,
        color,
        value: node.value
    };
}

function tryChangeColor(node: RedBlackTree<any>, color: Color): RedBlackTree<any> {
    if (isLeaf(node)) {
        return Nil;
    } else {
        return {
            left: node.left,
            right: node.right,
            color,
            value: node.value
        };
    }
}

function isLeaf(tree: RedBlackTree<any>): tree is RedBlackTreeLeaf {
    return tree === null;
}

function isNode<T>(tree: RedBlackTree<T>): tree is RedBlackTreeNode<T> {
    return tree !== null;
}

function isRedNode<T>(tree: RedBlackTree<T>): tree is RedBlackTreeNode<T> {
    return tree !== null && tree.color === Color.RED;
}

function isBlackNode<T>(tree: RedBlackTree<T>): tree is RedBlackTreeNode<T> {
    return tree !== null && tree.color === Color.BLACK;
}

function getBlackHeight(tree: RedBlackTree<any>): number {
    if (isLeaf(tree)) {
        return 0;
    } else {
        const leftHeight = getBlackHeight(tree.left);
        const rightHeight = getBlackHeight(tree.right);
        if (leftHeight < 0 || rightHeight < 0 || leftHeight !== rightHeight) {
            return -1;
        } else {
            return leftHeight + (tree.color === Color.BLACK ? 1 : 0);
        }
    }
}

function showTree(tree: RedBlackTree<any>): string {
    return showTreeStep(tree, 0, []).join('\n');
}

function showTreeStep(tree: RedBlackTree<any>, depth: number, output: string[]): string[] {
    const COLORS = { [Color.RED]: 'R', [Color.BLACK]: 'B' };
    if (isNode(tree)) {
        showTreeStep(tree.right, depth + 1, output);
        output.push('\t'.repeat(depth) + COLORS[tree.color] + ' ' + tree.value);
        showTreeStep(tree.left, depth + 1, output);
    }
    return output;
}