Browse Source

rbtree tests complete

tags/before-insert-rework
Jim Paris 9 years ago
parent
commit
c396c4dac8
2 changed files with 21 additions and 20 deletions
  1. +10
    -16
      nilmdb/rbtree.py
  2. +11
    -4
      tests/test_rbtree.py

+ 10
- 16
nilmdb/rbtree.py View File

@@ -358,19 +358,13 @@ class RBTree(object):
def intersect(self, start, end):
"""Generator that returns nodes that overlap the given
(start,end) range. Assumes non-overlapping intervals."""

# Start with the leftmost node before the starting point
n = self.find_left(start, start)
# If we didn't find one, look for the leftmode node before the
# ending point instead.
if n is None:
n = self.find_left(end, end)
# If we still didn't find it, there are no intervals that intersect.
if n is None:
return

# Now yield this node and all successors until their endpoints

if False:
yield
return
# Start with the leftmode node that ends after start
n = self.find_left_end(start)
while n is not None:
if n.start >= end:
# this node starts after the requested end; we're done
break
if start < n.end:
# this node overlaps our requested area
yield n
n = self.successor(n)

+ 11
- 4
tests/test_rbtree.py View File

@@ -156,8 +156,15 @@ class TestRBTree:
assert(node.start > i)

def test_rbtree_intersect(self):
# Fill with some ranges again
# Fill with some ranges
rb = RBTree()
vals = [ 7, 14, 1, 2, 8, 11, 5, 15, 4]
for n in vals:
rb.insert(RBNode(n*10, n*10+5))
rb.insert(RBNode(10,20))
rb.insert(RBNode(20,25))
rb.insert(RBNode(30,40))
# Just a quick test; test_interval will do better.
eq_(len(list(rb.intersect(1,100))), 3)
eq_(len(list(rb.intersect(10,20))), 1)
eq_(len(list(rb.intersect(5,15))), 1)
eq_(len(list(rb.intersect(15,15))), 1)
eq_(len(list(rb.intersect(20,21))), 1)
eq_(len(list(rb.intersect(19,21))), 2)

Loading…
Cancel
Save