diff --git a/Readme.md b/Readme.md index 2e51d6d..f894299 100644 --- a/Readme.md +++ b/Readme.md @@ -53,6 +53,8 @@ Or use the compiled version 'dist/splay.js'. - `tree.remove(key:any)` - Remove item - `tree.find(key):Node|Null` - Return node by its key - `tree.findStatic(key):Node|Null` - Return node by its key (doesn't re-balance the tree) +- `tree.findAtOrNeighbor(key):Node|Null` - Return node by its key or a node with a key right above or below. Useful for floating point keys, etc. +- `tree.findAtOrNeighborStatic(key):Node|Null` - Same as above, but doesn't re-balance the tree - `tree.at(index:Number):Node|Null` - Return node by its index in sorted order of keys - `tree.contains(key):Boolean` - Whether a node with the given key is in the tree - `tree.forEach(function(node) {...}):Tree` In-order traversal diff --git a/src/index.ts b/src/index.ts index 4c04c1d..2f61dcd 100644 --- a/src/index.ts +++ b/src/index.ts @@ -260,31 +260,49 @@ export default class Tree { return null; } + public findAtOrNeighborStatic (key:Key) : Node|null { + let current = this._root; + const compare = this._comparator; + let next:Node; + while (current) { + const cmp = compare(key, current.key); + if (cmp === 0) return current; + else if (cmp < 0) next = current.left; + else next = current.right; + + if (next == undefined) return current; + + current = next; + } + } + /** * Find without splaying */ public findStatic (key:Key) : Node|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; + const result_candidate = this.findAtOrNeighborStatic(key); + if (result_candidate && this._comparator(result_candidate.key, key) === 0) { + return result_candidate; } return null; } - - public find (key:Key) : Node|null { + public findAtOrNeighbor (key:Key) : Node|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 find (key:Key) : Node|null { + const result_candidate = this.findAtOrNeighbor(key); + if (result_candidate && this._comparator(result_candidate.key, key) === 0) { + return result_candidate; + } + return null; + } + public contains (key:Key) : boolean { let current = this._root; diff --git a/tests/find.test.ts b/tests/find.test.ts index eccd9dd..f62e810 100644 --- a/tests/find.test.ts +++ b/tests/find.test.ts @@ -56,4 +56,66 @@ describe ('find', () => { assert.strictEqual(tree.find(2), tree.root); }); + + it('should return previous key as result of search', () => { + const tree = new Tree(); + assert.equal(tree.findAtOrNeighbor(10), null); + assert.equal(tree.findAtOrNeighbor(20), null); + assert.equal(tree.findAtOrNeighbor(30), null); + tree.insert(10, 40); + tree.insert(20, 50); + tree.insert(30, 60); + + let root = tree.root; + assert.equal(tree.findAtOrNeighbor(10).data, 40); + assert.notStrictEqual(root, tree.root); + root = tree.root; + + const found_around_15 = tree.findAtOrNeighbor(15); + assert.equal(found_around_15.data === 40 || found_around_15.data === 50, true); + + assert.equal(tree.findAtOrNeighbor(20).data, 50); + assert.notStrictEqual(root, tree.root); + root = tree.root; + + assert.equal(tree.findAtOrNeighbor(30).data, 60); + assert.notStrictEqual(root, tree.root); + root = tree.root; + + assert.equal(tree.findAtOrNeighbor(80).data, 60); + assert.strictEqual(root, tree.root); + }); + + it ('should allow finding node by approximate key without splaying', () => { + const tree = new Tree(); + assert.equal(tree.findAtOrNeighborStatic(10), null); + assert.equal(tree.findAtOrNeighborStatic(20), null); + assert.equal(tree.findAtOrNeighborStatic(30), null); + tree.insert(-20, 80); + tree.insert(10, 40); + tree.insert(20, 50); + tree.insert(30, 60); + + tree.find(20); + const root = tree.root; + assert.equal(tree.findAtOrNeighborStatic(10).data, 40); + assert.strictEqual(root, tree.root); + + const result_around_15 = tree.findAtOrNeighborStatic(15); + assert.equal(result_around_15.data === 40 || result_around_15.data === 50, true); + assert.strictEqual(root, tree.root); + + assert.equal(tree.findAtOrNeighborStatic(20).data, 50); + assert.strictEqual(root, tree.root); + + assert.equal(tree.findAtOrNeighborStatic(30).data, 60); + assert.strictEqual(root, tree.root); + + assert.equal(tree.findAtOrNeighborStatic(80).data, 60); + + assert.equal(tree.findAtOrNeighborStatic(-20).data, 80); + assert.equal(tree.findAtOrNeighborStatic(-30).data, 80); + + assert.strictEqual(tree.find(20), tree.root); + }); });