Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add methods first_before and first_after. #39

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 53 additions & 0 deletions intervaltree/intervaltree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
45 changes: 45 additions & 0 deletions test/intervaltree_methods/first_before_after_test.py
Original file line number Diff line number Diff line change
@@ -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'])