diff --git a/intervaltree/intervaltree.py b/intervaltree/intervaltree.py index 3eda476..955bb46 100644 --- a/intervaltree/intervaltree.py +++ b/intervaltree/intervaltree.py @@ -762,6 +762,59 @@ def items(self): :rtype: set of Interval """ return set(self.all_intervals) + + def _first_after_or_after(self, value, key, optimum, compare): + node = self.top_node + first = None + if not node: + raise ValueError('Empty IntervalTree.') + + def update_first(first, ivs): + try: + opt = optimum((iv for iv in ivs if compare(iv.begin, value)), key=key) + except ValueError: + pass + else: + if first is None or compare(key(first), key(opt)): + first = opt + return first + + if node.x_center < value: + while node is not None and not compare(node.x_center, value): + first = update_first(first, node.s_center) + node = node.right_node + else: + while node is not None and compare(node.x_center, value): + first = update_first(first, node.s_center) + node = node.left_node + if node is None: + if first: + return first + else: + raise ValueError('No interval before/after {0}.'.format(value)) + else: + candidates = [iv for iv in node.all_children() if compare(key(iv), value)] + if first: + candidates.append(first) + if not candidates: + raise ValueError('No interval before/after {0}.'.format(value)) + return optimum(candidates, key=key) + + def first_after(self, value): + """Return an interval whose beginning is ≥ value and is minimal.""" + return self._first_after_or_after(value, + key=lambda iv: iv.begin, + optimum=min, + compare=lambda x, y: x >= y) + + def first_before(self, value): + """Return an interval whose end is ≤ value and is maximal.""" + return self._first_after_or_after(value, + key=lambda iv: iv.end, + optimum=max, + compare=lambda x, y: x <= y) + + def is_empty(self): """ diff --git a/test/intervaltree_methods/first_before_after_test.py b/test/intervaltree_methods/first_before_after_test.py new file mode 100644 index 0000000..2c273b8 --- /dev/null +++ b/test/intervaltree_methods/first_before_after_test.py @@ -0,0 +1,45 @@ +from __future__ import absolute_import +from intervaltree import Interval, IntervalTree +import pytest +from test.intervaltrees import trees, sdata + +def test_first_after(): + t = trees['ivs1']() + assert t.first_after(7) in [ + Interval(8, 10, '[8,10)'), + Interval(8, 15, '[8,15)'), + ] + assert t.first_after(8) in [ + Interval(8, 10, '[8,10)'), + Interval(8, 15, '[8,15)'), + ] + assert t.first_after(13) == Interval(14, 15, '[14,15)') + assert t.first_after(4) == Interval(4, 7, '[4,7)') + assert t.first_after(5) == Interval(5, 9, '[5,9)') + assert t.first_after(-100) == Interval(1, 2, '[1,2)') + try: + iv = t.first_after(15) + except Exception as e: + assert isinstance(e, ValueError) + else: + assert False, iv + +def test_first_before(): + t = trees['ivs1']() + assert t.first_before(5) == Interval(1, 2, '[1,2)') + assert t.first_before(7) == Interval(4, 7, '[4,7)') + assert t.first_before(10) in [ + Interval(6, 10, '[6,10)'), + Interval(8, 10, '[8,10)'), + ] + assert t.first_before(15) == Interval(14, 15, '[14,15)') + assert t.first_before(100) == Interval(14, 15, '[14,15)') + try: + iv = t.first_before(1) + except Exception as e: + assert isinstance(e, ValueError) + else: + assert False, iv + +if __name__ == "__main__": + pytest.main([__file__, '-v'])