diff --git a/README.md b/README.md index 947fe0e..60d7e81 100644 --- a/README.md +++ b/README.md @@ -132,11 +132,25 @@ Returns an array of data items (points or rectangles) that the given bounding bo Note that the `search` method accepts a bounding box in `{minX, minY, maxX, maxY}` format regardless of the format specified in the constructor (which only affects inserted objects). +Also `search` accepts a predicate, which is a function that takes a single result as an argument and returns `true` or `false` if the item should be included to the final result or not respectively: ```js -var allItems = tree.all(); +var result = tree.search({ + minX: 40, + minY: 20, + maxX: 80, + maxY: 70 +}, function (dot) { + if (dot.someCheckHere) { + return false; + } + return true; +}); ``` -Returns all items of the tree. +Return all items of the tree: +```js +var allItems = tree.all(); +``` ### Collisions diff --git a/index.js b/index.js index 41c3b7e..b6b9f14 100644 --- a/index.js +++ b/index.js @@ -24,7 +24,7 @@ rbush.prototype = { return this._all(this.data, []); }, - search: function (bbox) { + search: function (bbox, predicate) { var node = this.data, result = [], @@ -42,8 +42,15 @@ rbush.prototype = { childBBox = node.leaf ? toBBox(child) : child; if (intersects(bbox, childBBox)) { - if (node.leaf) result.push(child); - else if (contains(bbox, childBBox)) this._all(child, result); + if (node.leaf) { + if (predicate) { + if (predicate(child)) { + result.push(child); + } + } else { + result.push(child); + } + } else if (contains(bbox, childBBox)) this._all(child, result, predicate); else nodesToSearch.push(child); } } @@ -187,11 +194,21 @@ rbush.prototype = { return this; }, - _all: function (node, result) { + _all: function (node, result, predicate) { var nodesToSearch = []; while (node) { - if (node.leaf) result.push.apply(result, node.children); - else nodesToSearch.push.apply(nodesToSearch, node.children); + if (node.leaf) { + if (predicate) { + for (var i = 0; i < node.children.length; i++) { + var child = node.children[i]; + if (predicate(child)) { + result.push(child); + } + } + } else { + result.push.apply(result, node.children); + } + } else nodesToSearch.push.apply(nodesToSearch, node.children); node = nodesToSearch.pop(); } diff --git a/test/test.js b/test/test.js index 21d6a44..0e51ee4 100644 --- a/test/test.js +++ b/test/test.js @@ -367,6 +367,36 @@ t('#clear should clear all the data in the tree', function (t) { t.end(); }); +t('#search accepts predicate: exclude results', function (t) { + var tree = rbush(4).load(data); + var predicate = function (point) { + if (point.isReturned) { + return false; + } + point.isReturned = true; + return true; + }; + var result = tree.search({minX: 40, minY: 20, maxX: 80, maxY: 70}, predicate); + + var setIsReturned = function (point) { + point.isReturned = true; + return point; + }; + + sortedEqual(t, result, [ + [70,20,70,20],[75,25,75,25],[45,45,45,45],[50,50,50,50],[60,60,60,60],[70,70,70,70], + [45,20,45,20],[45,70,45,70],[75,50,75,50],[50,25,50,25],[60,35,60,35],[70,45,70,45] + ].map(arrToBBox).map(setIsReturned)); + + result = tree.search({minX: 35, minY: 20, maxX: 80, maxY: 70}, predicate); + + t.equal(result.length, 2); + sortedEqual(t, result, [ + [35, 35, 35, 35], [35, 60, 35, 60] + ].map(arrToBBox).map(setIsReturned)); + t.end(); +}); + t('should have chainable API', function (t) { t.doesNotThrow(function () { rbush()