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         if low > high:
  26             temp: Numeric = low
  27             low = high
  28             high = temp
  29         self.low: Numeric = low
  30         self.high: Numeric = high
  31         self.highest_in_subtree: Numeric = high
  32
  33     def __lt__(self, other: NumericRange) -> bool:
  34         return self.low < other.low
  35
  36     @overrides
  37     def __eq__(self, other: object) -> bool:
  38         if not isinstance(other, NumericRange):
  39             return False
  40         return self.low == other.low and self.high == other.high
  41
  42     def overlaps_with(self, other: NumericRange) -> bool:
  43         return self.low <= other.high and self.high >= other.low
  44
  45     def __repr__(self) -> str:
  46         return f"{self.low}..{self.high}"
  47
  48
  49 class AugmentedIntervalTree(bst.BinarySearchTree):
  50     @staticmethod
  51     def _assert_value_must_be_range(value: Any) -> None:
  52         if not isinstance(value, NumericRange):
  53             raise Exception(
  54                 "AugmentedIntervalTree expects to use NumericRanges, see bst for a "
  55                 + "general purpose tree usable for other types."
  56             )
  57
  58     @overrides
  59     def _on_insert(self, parent: Optional[bst.Node], new: bst.Node) -> None:
  60         AugmentedIntervalTree._assert_value_must_be_range(new.value)
  61         for ancestor in self.parent_path(new):
  62             assert ancestor
  63             if new.value.high > ancestor.value.highest_in_subtree:
  64                 ancestor.value.highest_in_subtree = new.value.high
  65
  66     @overrides
  67     def _on_delete(self, parent: Optional[bst.Node], deleted: bst.Node) -> None:
  68         if parent:
  69             new_highest_candidates = [parent.value.high]
  70             if parent.left:
  71                 new_highest_candidates.append(parent.left.value.highest_in_subtree)
  72             if parent.right:
  73                 new_highest_candidates.append(parent.right.value.highest_in_subtree)
  74             parent.value.highest_in_subtree = max(new_highest_candidates)
  75
  76     def find_one_overlap(self, x: NumericRange):
  77         """Identify and return one overlapping node from the tree.
  78
  79         >>> tree = AugmentedIntervalTree()
  80         >>> tree.insert(NumericRange(20, 24))
  81         >>> tree.insert(NumericRange(18, 22))
  82         >>> tree.insert(NumericRange(14, 16))
  83         >>> tree.insert(NumericRange(1, 30))
  84         >>> tree.insert(NumericRange(25, 30))
  85         >>> tree.insert(NumericRange(29, 33))
  86         >>> tree.insert(NumericRange(5, 12))
  87         >>> tree.insert(NumericRange(1, 6))
  88         >>> tree.insert(NumericRange(13, 18))
  89         >>> tree.insert(NumericRange(16, 28))
  90         >>> tree.insert(NumericRange(21, 27))
  91         >>> tree.find_one_overlap(NumericRange(6, 7))
  92         1..30
  93         """
  94         return self._find_one_overlap(self.root, x)
  95
  96     def _find_one_overlap(
  97         self, root: bst.Node, x: NumericRange
  98     ) -> Optional[NumericRange]:
  99         if root is None:
 100             return None
 101
 102         if root.value.overlaps_with(x):
 103             return root.value
 104
 105         if root.left:
 106             if root.left.value.highest_in_subtree >= x.low:
 107                 return self._find_one_overlap(root.left, x)
 108
 109         if root.right:
 110             return self._find_one_overlap(root.right, x)
 111         return None
 112
 113     def find_all_overlaps(self, x: NumericRange):
 114         """Yields ranges previously added to the tree that x overlaps with.
 115
 116         >>> tree = AugmentedIntervalTree()
 117         >>> tree.insert(NumericRange(20, 24))
 118         >>> tree.insert(NumericRange(18, 22))
 119         >>> tree.insert(NumericRange(14, 16))
 120         >>> tree.insert(NumericRange(1, 30))
 121         >>> tree.insert(NumericRange(25, 30))
 122         >>> tree.insert(NumericRange(29, 33))
 123         >>> tree.insert(NumericRange(5, 12))
 124         >>> tree.insert(NumericRange(1, 6))
 125         >>> tree.insert(NumericRange(13, 18))
 126         >>> tree.insert(NumericRange(16, 28))
 127         >>> tree.insert(NumericRange(21, 27))
 128         >>> for x in tree.find_all_overlaps(NumericRange(19, 21)):
 129         ...     print(x)
 130         20..24
 131         18..22
 132         1..30
 133         16..28
 134         21..27
 135
 136         >>> del tree[NumericRange(1, 30)]
 137         >>> for x in tree.find_all_overlaps(NumericRange(19, 21)):
 138         ...     print(x)
 139         20..24
 140         18..22
 141         16..28
 142         21..27
 143         """
 144         if self.root is None:
 145             return
 146         yield from self._find_all_overlaps(self.root, x)
 147
 148     def _find_all_overlaps(
 149         self, root: bst.Node, x: NumericRange
 150     ) -> Generator[NumericRange, None, None]:
 151         if root is None:
 152             return None
 153
 154         if root.value.overlaps_with(x):
 155             yield root.value
 156
 157         if root.left:
 158             if root.left.value.highest_in_subtree >= x.low:
 159                 yield from self._find_all_overlaps(root.left, x)
 160
 161         if root.right:
 162             if root.right.value.highest_in_subtree >= x.low:
 163                 yield from self._find_all_overlaps(root.right, x)
 164
 165
 166 if __name__ == "__main__":
 167     import doctest
 168
 169     doctest.testmod()