diff --git a/CHANGELOG.md b/CHANGELOG.md index a71113f..4791347 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,9 @@ # Change log ## Version 3.2.2 +- Added: + - `IntervalTree.nearest(target, k=None)` finds the interval closest to a point or + Interval, or the `k` closest intervals, nearest first. - Fixed: - Better subclassing support: Determine classes dynamically, so that methods like str() are aware when our types are subclassed. diff --git a/README.md b/README.md index a9635c7..6aa983e 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,11 @@ Features * Envelop queries * `tree.envelop(begin, end)` +* Nearest queries + * `tree.nearest(point)` (the closest `Interval`, or `None` if the tree is empty) + * `tree.nearest(interval)` + * `tree.nearest(point, k=3)` (a list of the 3 closest intervals, nearest first) + * Membership queries * `interval_obj in tree` (this is fastest, O(1)) * `tree.containsi(begin, end, data)` diff --git a/intervaltree/intervaltree.py b/intervaltree/intervaltree.py index e273127..94d65cf 100644 --- a/intervaltree/intervaltree.py +++ b/intervaltree/intervaltree.py @@ -24,7 +24,7 @@ """ from .interval import Interval from .node import Node -from numbers import Number +from numbers import Integral, Number from sortedcontainers import SortedDict from copy import copy @@ -944,6 +944,75 @@ def overlap(self, begin, end=None): )) return result + def nearest(self, target, k=None): + """ + Returns the interval closest to target, or the k closest intervals. + + The target is a point or an Interval. The distance to an interval + is given by Interval.distance_to(): 0 when the two overlap or touch, + otherwise the size of the gap between them. Ties are broken by the + ordering of the intervals themselves (begin, then end, then data). + + If k is None, returns a single Interval, or None if the tree is + empty. Otherwise returns a list of at most k intervals, nearest + first. + + Completes in O(m + log n) time when the neighbours are close to the + target, where: + * n = size of the tree + * m = number of intervals in the search window, which grows + until it holds at least k intervals + + :param target: a point or an Interval + :param k: None, or the maximum number of intervals to return + :raises ValueError: if target is a null Interval, or k is negative + :raises TypeError: if k is neither None nor an integer + :rtype: Interval, None, or list of Interval + """ + if isinstance(target, Interval): + if target.is_null(): + raise ValueError( + "Cannot search near a null Interval: {0}".format(target) + ) + begin, end = target.begin, target.end + else: + begin = end = target + + if k is None: + found = self.nearest(target, 1) + return found[0] if found else None + if isinstance(k, bool) or not isinstance(k, Integral): + raise TypeError("k must be an integer or None, not {0!r}".format(k)) + if k < 0: + raise ValueError("k must not be negative, not {0!r}".format(k)) + if k == 0 or not self: + return [] + + # Every interval within `radius` of the target has an endpoint in + # [begin - radius, end + radius], or spans that whole range. Since + # all endpoints are in the boundary table, widening the search to the + # neighbouring boundaries turns overlap()'s half-open test into a + # closed one, so intervals exactly `radius` away are included too. + table = self.boundary_table + keys = table.keys() + radius = 0 + while True: + i = table.bisect_left(begin - radius) + j = table.bisect_right(end + radius) + lower = keys[i - 1] if i else keys[0] + upper = keys[j] if j < len(keys) else keys[-1] + found = self.overlap(lower, upper) + if len(found) >= k or len(found) == len(self): + break + gaps = [] + if i: + gaps.append(begin - keys[i - 1]) + if j < len(keys): + gaps.append(keys[j] - end) + radius = max(2 * radius, min(gaps)) + found = sorted(found, key=lambda iv: (iv.distance_to(target), iv)) + return found[:k] + def begin(self): """ Returns the lower bound of the first interval in the tree. diff --git a/test/intervaltree_methods/nearest_test.py b/test/intervaltree_methods/nearest_test.py new file mode 100644 index 0000000..f6ae391 --- /dev/null +++ b/test/intervaltree_methods/nearest_test.py @@ -0,0 +1,168 @@ +""" +intervaltree: A mutable, self-balancing interval tree for Python 2 and 3. +Queries may be by point, by range overlap, or by range envelopment. + +Test module: IntervalTree, nearest() queries + +Copyright 2013-2018 Chaim Leib Halbert + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" +from __future__ import absolute_import +from intervaltree import Interval, IntervalTree +import pytest +from random import Random + + +def brute_force(intervals, target, k): + ranked = sorted(intervals, key=lambda iv: (iv.distance_to(target), iv)) + return ranked[:k] + + +def test_empty_tree(): + t = IntervalTree() + assert t.nearest(5) is None + assert t.nearest(5, k=3) == [] + assert t.nearest(Interval(1, 2)) is None + + +def test_point_inside_interval(): + t = IntervalTree.from_tuples([(0, 5, 'a'), (10, 15, 'b')]) + assert t.nearest(3) == Interval(0, 5, 'a') + assert t.nearest(10) == Interval(10, 15, 'b') + + +def test_point_in_gap(): + t = IntervalTree.from_tuples([(0, 5, 'a'), (10, 15, 'b')]) + assert t.nearest(6) == Interval(0, 5, 'a') + assert t.nearest(8) == Interval(10, 15, 'b') + assert t.nearest(-100) == Interval(0, 5, 'a') + assert t.nearest(100) == Interval(10, 15, 'b') + + +def test_point_at_open_end(): + t = IntervalTree.from_tuples([(0, 5, 'a'), (6, 9, 'b')]) + assert t.nearest(5, k=2) == [Interval(0, 5, 'a'), Interval(6, 9, 'b')] + assert Interval(0, 5, 'a').distance_to(5) == 0 + + +def test_interval_target(): + t = IntervalTree.from_tuples([(0, 5), (10, 15), (30, 40)]) + assert t.nearest(Interval(4, 8)) == Interval(0, 5) + assert t.nearest(Interval(16, 20)) == Interval(10, 15) + assert t.nearest(Interval(22, 26)) == Interval(30, 40) + assert t.nearest(Interval(-50, 100)) == Interval(0, 5) + + +def test_touching_interval_is_at_distance_zero(): + t = IntervalTree.from_tuples([(0, 5), (7, 9)]) + assert t.nearest(Interval(5, 6), k=1) == [Interval(0, 5)] + assert t.nearest(Interval(6, 7), k=2) == [Interval(7, 9), Interval(0, 5)] + + +def test_k_returns_list_nearest_first(): + t = IntervalTree.from_tuples([(0, 1), (3, 4), (10, 11), (20, 21)]) + assert t.nearest(5, k=3) == [Interval(3, 4), Interval(0, 1), Interval(10, 11)] + assert t.nearest(5, k=1) == [Interval(3, 4)] + + +def test_k_zero(): + t = IntervalTree.from_tuples([(0, 1)]) + assert t.nearest(0, k=0) == [] + + +def test_k_larger_than_tree(): + t = IntervalTree.from_tuples([(0, 1), (5, 6), (9, 12)]) + assert t.nearest(4, k=10) == [Interval(5, 6), Interval(0, 1), Interval(9, 12)] + + +def test_ties_are_broken_by_interval_order(): + t = IntervalTree.from_tuples([(10, 12, 'x'), (0, 4, 'y'), (0, 4, 'x')]) + assert t.nearest(7, k=3) == [ + Interval(0, 4, 'x'), + Interval(0, 4, 'y'), + Interval(10, 12, 'x'), + ] + assert t.nearest(7) == Interval(0, 4, 'x') + + +def test_ties_with_unorderable_data(): + t = IntervalTree([Interval(0, 4, 'a'), Interval(0, 4, 1), Interval(10, 12)]) + result = t.nearest(7, k=3) + assert result[:2] == sorted([Interval(0, 4, 'a'), Interval(0, 4, 1)]) + assert result[2] == Interval(10, 12) + + +def test_float_and_negative_coordinates(): + t = IntervalTree.from_tuples([(-2.5, -1.5), (0.5, 1.0), (4.25, 9.0)]) + assert t.nearest(-1.0) == Interval(-2.5, -1.5) + assert t.nearest(2.0) == Interval(0.5, 1.0) + assert t.nearest(3.0, k=2) == [Interval(4.25, 9.0), Interval(0.5, 1.0)] + + +def test_nearest_after_mutation(): + t = IntervalTree.from_tuples([(0, 1), (100, 101)]) + assert t.nearest(50) == Interval(0, 1) + t.addi(48, 49) + assert t.nearest(50) == Interval(48, 49) + t.remove(Interval(48, 49)) + assert t.nearest(50) == Interval(0, 1) + + +def test_tree_is_not_modified(): + t = IntervalTree.from_tuples([(0, 1), (5, 6), (9, 12)]) + before = set(t) + t.nearest(4, k=2) + assert set(t) == before + t.verify() + + +def test_null_interval_target(): + t = IntervalTree.from_tuples([(0, 1)]) + with pytest.raises(ValueError): + t.nearest(Interval(3, 3)) + with pytest.raises(ValueError): + t.nearest(Interval(5, 3), k=2) + + +def test_invalid_k(): + t = IntervalTree.from_tuples([(0, 1)]) + with pytest.raises(ValueError): + t.nearest(0, k=-1) + for bad in (1.5, '2', True): + with pytest.raises(TypeError): + t.nearest(0, k=bad) + + +def test_against_brute_force(): + rng = Random(20240229) + for _ in range(300): + intervals = set() + for _ in range(rng.randint(0, 25)): + begin = rng.randint(-20, 20) + intervals.add(Interval(begin, begin + rng.randint(1, 8), + rng.choice([None, 'a', 'b']))) + t = IntervalTree(intervals) + if rng.random() < 0.5: + target = rng.randint(-30, 30) + else: + begin = rng.randint(-30, 30) + target = Interval(begin, begin + rng.randint(1, 8)) + for k in (1, 2, 3, 7, 40): + assert t.nearest(target, k=k) == brute_force(intervals, target, k) + expected = brute_force(intervals, target, 1) + assert t.nearest(target) == (expected[0] if expected else None) + + +if __name__ == "__main__": + pytest.main([__file__, '-v'])