src/avlTree.ts - vscode-go - Git at Google

 /* eslint-disable @typescript-eslint/no-explicit-any */
 /**
  * @license
  * Copyright Daniel Imms <http://www.growingwiththeweb.com>
  * Released under MIT license:
  *
  * The MIT License (MIT)
  *
  * Copyright (c) 2016 Daniel Imms, http://www.growingwiththeweb.com
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
  * in the Software without restriction, including without limitation the rights
  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  * copies of the Software, and to permit persons to whom the Software is
  * furnished to do so, subject to the following conditions:
  *
  * The above copyright notice and this permission notice shall be included in
  * all copies or substantial portions of the Software.
  *
  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  *
  * Modified by Jackson Kearl <Microsoft/t-jakea@microsoft.com>
  */

 /**
  * Represents a node in the binary tree, which has a key and a value, as well as left and right subtrees
  */
 export class Node<K, V> {
 	public left: Node<K, V> | null = null;
 	public right: Node<K, V> | null = null;
 	public height = 0;

 	/**
 	 * Creates a new AVL Tree node.
 	 * @param key The key of the new node.
 	 * @param value The value of the new node.
 	 */
 	constructor(public key: K, public value: V | undefined) {}

 	/**
 	 * Convenience function to get the height of the left child of the node,
 	 * returning -1 if the node is null.
 	 * @return The height of the left child, or -1 if it doesn't exist.
 	 */
 	public get leftHeight(): number {
 		if (!this.left) {
 			return -1;
 		}
 		return this.left.height;
 	}

 	/**
 	 * Convenience function to get the height of the right child of the node,
 	 * returning -1 if the node is null.
 	 * @return The height of the right child, or -1 if it doesn't exist.
 	 */
 	public get rightHeight(): number {
 		if (!this.right) {
 			return -1;
 		}
 		return this.right.height;
 	}

 	/**
 	 * Performs a right rotate on this node.
 	 * @return The root of the sub-tree; the node where this node used to be.
 	 */
 	public rotateRight(): Node<K, V> {
 		//     b                           a
 		//    / \                         / \
 		//   a   e -> b.rotateRight() -> c   b
 		//  / \                             / \
 		// c   d                           d   e
 		const other = <Node<K, V>>this.left;
 		this.left = other.right;
 		other.right = this;
 		this.height = Math.max(this.leftHeight, this.rightHeight) + 1;
 		other.height = Math.max(other.leftHeight, this.height) + 1;
 		return other;
 	}

 	/**
 	 * Performs a left rotate on this node.
 	 * @return The root of the sub-tree; the node where this node used to be.
 	 */
 	public rotateLeft(): Node<K, V> {
 		//   a                              b
 		//  / \                            / \
 		// c   b   -> a.rotateLeft() ->   a   e
 		//    / \                        / \
 		//   d   e                      c   d
 		const other = <Node<K, V>>this.right;
 		this.right = other.left;
 		other.left = this;
 		this.height = Math.max(this.leftHeight, this.rightHeight) + 1;
 		other.height = Math.max(other.rightHeight, this.height) + 1;
 		return other;
 	}
 }

 export type DistanceFunction<K> = (a: K, b: K) => number;
 export type CompareFunction<K> = (a: K, b: K) => number;

 /**
  * Represents how balanced a node's left and right children are.
  */
 const enum BalanceState {
 	/** Right child's height is 2+ greater than left child's height */
 	UNBALANCED_RIGHT,
 	/** Right child's height is 1 greater than left child's height */
 	SLIGHTLY_UNBALANCED_RIGHT,
 	/** Left and right children have the same height */
 	BALANCED,
 	/** Left child's height is 1 greater than right child's height */
 	SLIGHTLY_UNBALANCED_LEFT,
 	/** Left child's height is 2+ greater than right child's height */
 	UNBALANCED_LEFT
 }

 export class NearestNeighborDict<K, V> {
 	public static NUMERIC_DISTANCE_FUNCTION = (a: number, b: number) => (a > b ? a - b : b - a);
 	public static DEFAULT_COMPARE_FUNCTION = (a: any, b: any) => (a > b ? 1 : a < b ? -1 : 0);

 	protected root: Node<K, V> | null = null;

 	/**
 	 * Creates a new AVL Tree.
 	 */
 	constructor(
 		start: Node<K, V>,
 		private distance: DistanceFunction<K>,
 		private compare: CompareFunction<K> = NearestNeighborDict.DEFAULT_COMPARE_FUNCTION
 	) {
 		this.insert(start.key, start.value);
 	}

 	public height() {
 		return this.root ? this.root.height : 0;
 	}

 	/**
 	 * Inserts a new node with a specific key into the tree.
 	 * @param key The key being inserted.
 	 * @param value The value being inserted.
 	 */
 	public insert(key: K, value?: V): void {
 		this.root = this._insert(key, value, this.root);
 	}

 	/**
 	 * Gets a node within the tree with a specific key, or the nearest neighbor to that node if it does not exist.
 	 * @param key The key being searched for.
 	 * @return The (key, value) pair of the node with key nearest the given key in value.
 	 */
 	public getNearest(key: K): Node<K, V> {
 		return this._getNearest(key, this.root!, this.root!);
 	}

 	/**
 	 * Inserts a new node with a specific key into the tree.
 	 * @param key The key being inserted.
 	 * @param root The root of the tree to insert in.
 	 * @return The new tree root.
 	 */
 	private _insert(key: K, value: V | undefined, root: Node<K, V> | null): Node<K, V> {
 		// Perform regular BST insertion
 		if (root === null) {
 			return new Node(key, value);
 		}

 		if (this.compare(key, root.key) < 0) {
 			root.left = this._insert(key, value, root.left);
 		} else if (this.compare(key, root.key) > 0) {
 			root.right = this._insert(key, value, root.right);
 		} else {
 			return root;
 		}

 		// Update height and rebalance tree
 		root.height = Math.max(root.leftHeight, root.rightHeight) + 1;
 		const balanceState = this._getBalanceState(root);

 		if (balanceState === BalanceState.UNBALANCED_LEFT) {
 			if (this.compare(key, root.left!.key) < 0) {
 				// Left left case
 				root = root.rotateRight();
 			} else {
 				// Left right case
 				root.left = root.left!.rotateLeft();
 				return root.rotateRight();
 			}
 		}

 		if (balanceState === BalanceState.UNBALANCED_RIGHT) {
 			if (this.compare(key, root.right!.key) > 0) {
 				// Right right case
 				root = root.rotateLeft();
 			} else {
 				// Right left case
 				root.right = root.right!.rotateRight();
 				return root.rotateLeft();
 			}
 		}

 		return root;
 	}

 	/**
 	 * Gets a node within the tree with a specific key, or the node closest (as measured by this._distance)
 	 * to that node if the key is not present
 	 * @param key The key being searched for.
 	 * @param root The root of the tree to search in.
 	 * @param closest The current best estimate of the node closest to the node being searched for,
 	 * as measured by this._distance
 	 * @return The (key, value) pair of the node with key nearest the given key in value.
 	 */
 	private _getNearest(key: K, root: Node<K, V>, closest: Node<K, V>): Node<K, V> {
 		const result = this.compare(key, root.key);
 		if (result === 0) {
 			return root;
 		}

 		closest = this.distance(key, root.key) < this.distance(key, closest.key) ? root : closest;

 		if (result < 0) {
 			return root.left ? this._getNearest(key, root.left, closest) : closest;
 		} else {
 			return root.right ? this._getNearest(key, root.right, closest) : closest;
 		}
 	}

 	/**
 	 * Gets the balance state of a node, indicating whether the left or right
 	 * sub-trees are unbalanced.
 	 * @param node The node to get the difference from.
 	 * @return The BalanceState of the node.
 	 */
 	private _getBalanceState(node: Node<K, V>): BalanceState {
 		const heightDifference = node.leftHeight - node.rightHeight;
 		switch (heightDifference) {
 			case -2:
 				return BalanceState.UNBALANCED_RIGHT;
 			case -1:
 				return BalanceState.SLIGHTLY_UNBALANCED_RIGHT;
 			case 1:
 				return BalanceState.SLIGHTLY_UNBALANCED_LEFT;
 			case 2:
 				return BalanceState.UNBALANCED_LEFT;
 			case 0:
 				return BalanceState.BALANCED;
 			default: {
 				console.error('Internal error: Avl tree should never be more than two levels unbalanced');
 				if (heightDifference > 0) {
 					return BalanceState.UNBALANCED_LEFT;
 				}
 				return BalanceState.UNBALANCED_RIGHT; // heightDifference can't be 0
 			}
 		}
 	}
 }
	/* eslint-disable @typescript-eslint/no-explicit-any */
	/**
	* @license
	* Copyright Daniel Imms <http://www.growingwiththeweb.com>
	* Released under MIT license:
	*
	* The MIT License (MIT)
	*
	* Copyright (c) 2016 Daniel Imms, http://www.growingwiththeweb.com
	*
	* Permission is hereby granted, free of charge, to any person obtaining a copy
	* of this software and associated documentation files (the "Software"), to deal
	* in the Software without restriction, including without limitation the rights
	* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	* copies of the Software, and to permit persons to whom the Software is
	* furnished to do so, subject to the following conditions:
	*
	* The above copyright notice and this permission notice shall be included in
	* all copies or substantial portions of the Software.
	*
	* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	* THE SOFTWARE.
	*
	* Modified by Jackson Kearl <Microsoft/t-jakea@microsoft.com>
	*/

	/**
	* Represents a node in the binary tree, which has a key and a value, as well as left and right subtrees
	*/
	export class Node<K, V> {
	public left: Node<K, V> \| null = null;
	public right: Node<K, V> \| null = null;
	public height = 0;

	/**
	* Creates a new AVL Tree node.
	* @param key The key of the new node.
	* @param value The value of the new node.
	*/
	constructor(public key: K, public value: V \| undefined) {}

	/**
	* Convenience function to get the height of the left child of the node,
	* returning -1 if the node is null.
	* @return The height of the left child, or -1 if it doesn't exist.
	*/
	public get leftHeight(): number {
	if (!this.left) {
	return -1;
	}
	return this.left.height;
	}

	/**
	* Convenience function to get the height of the right child of the node,
	* returning -1 if the node is null.
	* @return The height of the right child, or -1 if it doesn't exist.
	*/
	public get rightHeight(): number {
	if (!this.right) {
	return -1;
	}
	return this.right.height;
	}

	/**
	* Performs a right rotate on this node.
	* @return The root of the sub-tree; the node where this node used to be.
	*/
	public rotateRight(): Node<K, V> {
	// b a
	// / \ / \
	// a e -> b.rotateRight() -> c b
	// / \ / \
	// c d d e
	const other = <Node<K, V>>this.left;
	this.left = other.right;
	other.right = this;
	this.height = Math.max(this.leftHeight, this.rightHeight) + 1;
	other.height = Math.max(other.leftHeight, this.height) + 1;
	return other;
	}

	/**
	* Performs a left rotate on this node.
	* @return The root of the sub-tree; the node where this node used to be.
	*/
	public rotateLeft(): Node<K, V> {
	// a b
	// / \ / \
	// c b -> a.rotateLeft() -> a e
	// / \ / \
	// d e c d
	const other = <Node<K, V>>this.right;
	this.right = other.left;
	other.left = this;
	this.height = Math.max(this.leftHeight, this.rightHeight) + 1;
	other.height = Math.max(other.rightHeight, this.height) + 1;
	return other;
	}
	}

	export type DistanceFunction<K> = (a: K, b: K) => number;
	export type CompareFunction<K> = (a: K, b: K) => number;

	/**
	* Represents how balanced a node's left and right children are.
	*/
	const enum BalanceState {
	/** Right child's height is 2+ greater than left child's height */
	UNBALANCED_RIGHT,
	/** Right child's height is 1 greater than left child's height */
	SLIGHTLY_UNBALANCED_RIGHT,
	/** Left and right children have the same height */
	BALANCED,
	/** Left child's height is 1 greater than right child's height */
	SLIGHTLY_UNBALANCED_LEFT,
	/** Left child's height is 2+ greater than right child's height */
	UNBALANCED_LEFT
	}

	export class NearestNeighborDict<K, V> {
	public static NUMERIC_DISTANCE_FUNCTION = (a: number, b: number) => (a > b ? a - b : b - a);
	public static DEFAULT_COMPARE_FUNCTION = (a: any, b: any) => (a > b ? 1 : a < b ? -1 : 0);

	protected root: Node<K, V> \| null = null;

	/**
	* Creates a new AVL Tree.
	*/
	constructor(
	start: Node<K, V>,
	private distance: DistanceFunction<K>,
	private compare: CompareFunction<K> = NearestNeighborDict.DEFAULT_COMPARE_FUNCTION
	) {
	this.insert(start.key, start.value);
	}

	public height() {
	return this.root ? this.root.height : 0;
	}

	/**
	* Inserts a new node with a specific key into the tree.
	* @param key The key being inserted.
	* @param value The value being inserted.
	*/
	public insert(key: K, value?: V): void {
	this.root = this._insert(key, value, this.root);
	}

	/**
	* Gets a node within the tree with a specific key, or the nearest neighbor to that node if it does not exist.
	* @param key The key being searched for.
	* @return The (key, value) pair of the node with key nearest the given key in value.
	*/
	public getNearest(key: K): Node<K, V> {
	return this._getNearest(key, this.root!, this.root!);
	}

	/**
	* Inserts a new node with a specific key into the tree.
	* @param key The key being inserted.
	* @param root The root of the tree to insert in.
	* @return The new tree root.
	*/
	private _insert(key: K, value: V \| undefined, root: Node<K, V> \| null): Node<K, V> {
	// Perform regular BST insertion
	if (root === null) {
	return new Node(key, value);
	}

	if (this.compare(key, root.key) < 0) {
	root.left = this._insert(key, value, root.left);
	} else if (this.compare(key, root.key) > 0) {
	root.right = this._insert(key, value, root.right);
	} else {
	return root;
	}

	// Update height and rebalance tree
	root.height = Math.max(root.leftHeight, root.rightHeight) + 1;
	const balanceState = this._getBalanceState(root);

	if (balanceState === BalanceState.UNBALANCED_LEFT) {
	if (this.compare(key, root.left!.key) < 0) {
	// Left left case
	root = root.rotateRight();
	} else {
	// Left right case
	root.left = root.left!.rotateLeft();
	return root.rotateRight();
	}
	}

	if (balanceState === BalanceState.UNBALANCED_RIGHT) {
	if (this.compare(key, root.right!.key) > 0) {
	// Right right case
	root = root.rotateLeft();
	} else {
	// Right left case
	root.right = root.right!.rotateRight();
	return root.rotateLeft();
	}
	}

	return root;
	}

	/**
	* Gets a node within the tree with a specific key, or the node closest (as measured by this._distance)
	* to that node if the key is not present
	* @param key The key being searched for.
	* @param root The root of the tree to search in.
	* @param closest The current best estimate of the node closest to the node being searched for,
	* as measured by this._distance
	* @return The (key, value) pair of the node with key nearest the given key in value.
	*/
	private _getNearest(key: K, root: Node<K, V>, closest: Node<K, V>): Node<K, V> {
	const result = this.compare(key, root.key);
	if (result === 0) {
	return root;
	}

	closest = this.distance(key, root.key) < this.distance(key, closest.key) ? root : closest;

	if (result < 0) {
	return root.left ? this._getNearest(key, root.left, closest) : closest;
	} else {
	return root.right ? this._getNearest(key, root.right, closest) : closest;
	}
	}

	/**
	* Gets the balance state of a node, indicating whether the left or right
	* sub-trees are unbalanced.
	* @param node The node to get the difference from.
	* @return The BalanceState of the node.
	*/
	private _getBalanceState(node: Node<K, V>): BalanceState {
	const heightDifference = node.leftHeight - node.rightHeight;
	switch (heightDifference) {
	case -2:
	return BalanceState.UNBALANCED_RIGHT;
	case -1:
	return BalanceState.SLIGHTLY_UNBALANCED_RIGHT;
	case 1:
	return BalanceState.SLIGHTLY_UNBALANCED_LEFT;
	case 2:
	return BalanceState.UNBALANCED_LEFT;
	case 0:
	return BalanceState.BALANCED;
	default: {
	console.error('Internal error: Avl tree should never be more than two levels unbalanced');
	if (heightDifference > 0) {
	return BalanceState.UNBALANCED_LEFT;
	}
	return BalanceState.UNBALANCED_RIGHT; // heightDifference can't be 0
	}
	}
	}
	}