src/pyutils/collectionz/interval_tree.py

   1 #!/usr/bin/env python3
   2
   3 """This is an augmented interval tree for storing ranges and identifying overlaps as
   4 described by: https://en.wikipedia.org/wiki/Interval_tree.
   5 """
   6
   7 from __future__ import annotations
   8
   9 from functools import total_ordering
  10 from typing import Any, Generator, Optional, Union
  11
  12 from overrides import overrides
  13
  14 from pyutils.collectionz import bst
  15
  16 Numeric = Union[int, float]
  17
  18
  19 @total_ordering
  20 class NumericRange(object):
  21     """Essentially a tuple of numbers denoting a range with some added
  22     helper methods on it."""
  23
  24     def __init__(self, low: Numeric, high: Numeric):
  25         """Creates a NumericRange.
  26
  27         Args:
  28             low: the lowest point in the range (inclusive).
  29             high: the highest point in the range (inclusive).
  30
  31         .. warning::
  32
  33             If low > high this code swaps the parameters and keeps the range
  34             rather than raising.
  35
  36         """
  37         if low > high:
  38             temp: Numeric = low
  39             low = high
  40             high = temp
  41         self.low: Numeric = low
  42         self.high: Numeric = high
  43         self.highest_in_subtree: Numeric = high
  44
  45     def __lt__(self, other: NumericRange) -> bool:
  46         """
  47         Returns:
  48             True is this range is less than (lower low) other, else False.
  49         """
  50         return self.low < other.low
  51
  52     @overrides
  53     def __eq__(self, other: object) -> bool:
  54         """
  55         Returns:
  56             True if this is the same range as other, else False.
  57         """
  58         if not isinstance(other, NumericRange):
  59             return False
  60         return self.low == other.low and self.high == other.high
  61
  62     def overlaps_with(self, other: NumericRange) -> bool:
  63         """
  64         Returns:
  65             True if this NumericRange overlaps with other, else False.
  66         """
  67         return self.low <= other.high and self.high >= other.low
  68
  69     def __repr__(self) -> str:
  70         return f"{self.low}..{self.high}"
  71
  72
  73 class AugmentedIntervalTree(bst.BinarySearchTree):
  74     @staticmethod
  75     def _assert_value_must_be_range(value: Any) -> None:
  76         if not isinstance(value, NumericRange):
  77             raise Exception(
  78                 "AugmentedIntervalTree expects to use NumericRanges, see bst for a "
  79                 + "general purpose tree usable for other types."
  80             )
  81
  82     @overrides
  83     def _on_insert(self, parent: Optional[bst.Node], new: bst.Node) -> None:
  84         AugmentedIntervalTree._assert_value_must_be_range(new.value)
  85         for ancestor in self.parent_path(new):
  86             assert ancestor
  87             if new.value.high > ancestor.value.highest_in_subtree:
  88                 ancestor.value.highest_in_subtree = new.value.high
  89
  90     @overrides
  91     def _on_delete(self, parent: Optional[bst.Node], deleted: bst.Node) -> None:
  92         if parent:
  93             new_highest_candidates = [parent.value.high]
  94             if parent.left:
  95                 new_highest_candidates.append(parent.left.value.highest_in_subtree)
  96             if parent.right:
  97                 new_highest_candidates.append(parent.right.value.highest_in_subtree)
  98             parent.value.highest_in_subtree = max(new_highest_candidates)
  99
 100     def find_one_overlap(self, to_find: NumericRange) -> Optional[NumericRange]:
 101         """Identify and return one overlapping node from the tree.
 102
 103         Args:
 104             to_find: the interval with which to find an overlap.
 105
 106         Returns:
 107             An overlapping range from the tree or None if no such
 108             ranges exist in the tree at present.
 109
 110         >>> tree = AugmentedIntervalTree()
 111         >>> tree.insert(NumericRange(20, 24))
 112         >>> tree.insert(NumericRange(18, 22))
 113         >>> tree.insert(NumericRange(14, 16))
 114         >>> tree.insert(NumericRange(1, 30))
 115         >>> tree.insert(NumericRange(25, 30))
 116         >>> tree.insert(NumericRange(29, 33))
 117         >>> tree.insert(NumericRange(5, 12))
 118         >>> tree.insert(NumericRange(1, 6))
 119         >>> tree.insert(NumericRange(13, 18))
 120         >>> tree.insert(NumericRange(16, 28))
 121         >>> tree.insert(NumericRange(21, 27))
 122         >>> tree.find_one_overlap(NumericRange(6, 7))
 123         1..30
 124
 125         """
 126         return self._find_one_overlap(self.root, to_find)
 127
 128     def _find_one_overlap(
 129         self, root: bst.Node, x: NumericRange
 130     ) -> Optional[NumericRange]:
 131         if root is None:
 132             return None
 133
 134         if root.value.overlaps_with(x):
 135             return root.value
 136
 137         if root.left:
 138             if root.left.value.highest_in_subtree >= x.low:
 139                 return self._find_one_overlap(root.left, x)
 140
 141         if root.right:
 142             return self._find_one_overlap(root.right, x)
 143         return None
 144
 145     def find_all_overlaps(
 146         self, to_find: NumericRange
 147     ) -> Generator[NumericRange, None, None]:
 148         """Yields ranges previously added to the tree that overlaps with
 149         to_find argument.
 150
 151         Args:
 152             to_find: the interval with which to find all overlaps.
 153
 154         Returns:
 155             A (potentially empty) sequence of all ranges in the tree
 156             that overlap with the argument.
 157
 158         >>> tree = AugmentedIntervalTree()
 159         >>> tree.insert(NumericRange(20, 24))
 160         >>> tree.insert(NumericRange(18, 22))
 161         >>> tree.insert(NumericRange(14, 16))
 162         >>> tree.insert(NumericRange(1, 30))
 163         >>> tree.insert(NumericRange(25, 30))
 164         >>> tree.insert(NumericRange(29, 33))
 165         >>> tree.insert(NumericRange(5, 12))
 166         >>> tree.insert(NumericRange(1, 6))
 167         >>> tree.insert(NumericRange(13, 18))
 168         >>> tree.insert(NumericRange(16, 28))
 169         >>> tree.insert(NumericRange(21, 27))
 170         >>> for x in tree.find_all_overlaps(NumericRange(19, 21)):
 171         ...     print(x)
 172         20..24
 173         18..22
 174         1..30
 175         16..28
 176         21..27
 177
 178         >>> del tree[NumericRange(1, 30)]
 179         >>> for x in tree.find_all_overlaps(NumericRange(19, 21)):
 180         ...     print(x)
 181         20..24
 182         18..22
 183         16..28
 184         21..27
 185
 186         """
 187         if self.root is None:
 188             return
 189         yield from self._find_all_overlaps(self.root, to_find)
 190
 191     def _find_all_overlaps(
 192         self, root: bst.Node, x: NumericRange
 193     ) -> Generator[NumericRange, None, None]:
 194         if root is None:
 195             return None
 196
 197         if root.value.overlaps_with(x):
 198             yield root.value
 199
 200         if root.left:
 201             if root.left.value.highest_in_subtree >= x.low:
 202                 yield from self._find_all_overlaps(root.left, x)
 203
 204         if root.right:
 205             if root.right.value.highest_in_subtree >= x.low:
 206                 yield from self._find_all_overlaps(root.right, x)
 207
 208
 209 if __name__ == "__main__":
 210     import doctest
 211
 212     doctest.testmod()