Skip to content

Commit

Permalink
add optional equality comparator
Browse files Browse the repository at this point in the history
  • Loading branch information
noname0310 committed Jan 13, 2022
1 parent 589327d commit 6d4f9c0
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion tree.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,13 @@ export { Node }
export type LessOp<K> = (a: K, b: K) => boolean


/** Type for determining equality of instances:
* a function to return true if entry `a` is equals to entry `b`.
* @typeparam K key type, entries are sorted by key
*/
export type EqualOp<K> = (a: K, b: K) => boolean


/** A red-black tree written in TypeScript. The entries are stored sorted after
* the criterium `lessOp` passed tp the constructor or by default by the
* comparison operator `<` (less). Insertion, deletion and iteration have
Expand All @@ -24,6 +31,7 @@ export class Tree<K = string, V = any>implements Map<K, V> {
/** @internal */ _root: Node<K, V> = Node.nilNode
/** @internal */ _size: number = 0
/** @internal */ readonly _less: Less<K, Node<K, V>>
/** @internal */ readonly _equal: EqualOp<K>

/** Create a new red-black tree optionally with entries from `source` and
* the sorting criterium `lessOp`.
Expand All @@ -35,8 +43,10 @@ export class Tree<K = string, V = any>implements Map<K, V> {
constructor(
source?: IterableIterator<[K, V]>,
lessOp: LessOp<K> = (a, b) => a < b,
equalOp: EqualOp<K> = (a, b) => a === b,
) {
this._less = (a, b) => lessOp(a, b.key)
this._equal = equalOp
if (source)
for (const entry of source)
this._insertNode(new Node(entry[0], entry[1]))
Expand Down Expand Up @@ -199,7 +209,8 @@ export class Tree<K = string, V = any>implements Map<K, V> {
key: K,
node: Node<K, V> = this._root,
): Node<K, V> {
while (node.ok && key !== node.key)
const equalOp = this._equal
while (node.ok && !equalOp(key, node.key))
node = this._less(key, node) ? node._left : node._right
return node
}
Expand Down

0 comments on commit 6d4f9c0

Please sign in to comment.