Source code for tryalgo.interval_tree

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""\
Interval tree
christoph dürr - jill-jênn vie - 2013-2018
"""

from bisect import bisect_right


# snip{
# pylint: disable=too-many-arguments, too-few-public-methods
class _Node:
    def __init__(self, center, by_low, by_high, left, right):
        self.center = center
        self.by_low = by_low
        self.by_high = by_high
        self.left = left
        self.right = right


[docs]def interval_tree(intervals): """Construct an interval tree :param intervals: list of half-open intervals encoded as value pairs *[left, right)* :assumes: intervals are lexicographically ordered ``>>> assert intervals == sorted(intervals)`` :returns: the root of the interval tree :complexity: :math:`O(n)` """ if intervals == []: return None center = intervals[len(intervals) // 2][0] L = [] R = [] C = [] for I in intervals: if I[1] <= center: L.append(I) elif center < I[0]: R.append(I) else: C.append(I) by_low = sorted((I[0], I) for I in C) by_high = sorted((I[1], I) for I in C) IL = interval_tree(L) IR = interval_tree(R) return _Node(center, by_low, by_high, IL, IR)
[docs]def intervals_containing(t, p): """Query the interval tree :param t: root of the interval tree :param p: value :returns: a list of intervals containing p :complexity: O(log n + m), where n is the number of intervals in t, and m the length of the returned list """ INF = float('inf') if t is None: return [] if p < t.center: retval = intervals_containing(t.left, p) j = bisect_right(t.by_low, (p, (INF, INF))) for i in range(j): retval.append(t.by_low[i][1]) else: retval = intervals_containing(t.right, p) i = bisect_right(t.by_high, (p, (INF, INF))) for j in range(i, len(t.by_high)): retval.append(t.by_high[j][1]) return retval
# snip}