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, 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     def __init__(self, low: Numeric, high: Numeric):
  22         if low > high:
  23             temp: Numeric = low
  24             low = high
  25             high = temp
  26         self.low: Numeric = low
  27         self.high: Numeric = high
  28         self.highest_in_subtree: Numeric = high
  29
  30     def __lt__(self, other: NumericRange) -> bool:
  31         return self.low < other.low
  32
  33     @overrides
  34     def __eq__(self, other: object) -> bool:
  35         if not isinstance(other, NumericRange):
  36             return False
  37         return self.low == other.low and self.high == other.high
  38
  39     def overlaps_with(self, other: NumericRange) -> bool:
  40         return self.low <= other.high and self.high >= other.low
  41
  42     def __repr__(self) -> str:
  43         return f"{self.low}..{self.high}"
  44
  45
  46 class AugmentedIntervalTree(bst.BinarySearchTree):
  47     def __init__(self):
  48         super().__init__()
  49
  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(self, root: bst.Node, x: NumericRange):
  97         if root is None:
  98             return
  99
 100         if root.value.overlaps_with(x):
 101             return root.value
 102
 103         if root.left:
 104             if root.left.value.highest_in_subtree >= x.low:
 105                 return self._find_one_overlap(root.left, x)
 106
 107         if root.right:
 108             return self._find_one_overlap(root.right, x)
 109
 110     def find_all_overlaps(self, x: NumericRange):
 111         """Yields ranges previously added to the tree that x overlaps with.
 112
 113         >>> tree = AugmentedIntervalTree()
 114         >>> tree.insert(NumericRange(20, 24))
 115         >>> tree.insert(NumericRange(18, 22))
 116         >>> tree.insert(NumericRange(14, 16))
 117         >>> tree.insert(NumericRange(1, 30))
 118         >>> tree.insert(NumericRange(25, 30))
 119         >>> tree.insert(NumericRange(29, 33))
 120         >>> tree.insert(NumericRange(5, 12))
 121         >>> tree.insert(NumericRange(1, 6))
 122         >>> tree.insert(NumericRange(13, 18))
 123         >>> tree.insert(NumericRange(16, 28))
 124         >>> tree.insert(NumericRange(21, 27))
 125         >>> for x in tree.find_all_overlaps(NumericRange(19, 21)):
 126         ...     print(x)
 127         20..24
 128         18..22
 129         1..30
 130         16..28
 131         21..27
 132
 133         >>> del tree[NumericRange(1, 30)]
 134         >>> for x in tree.find_all_overlaps(NumericRange(19, 21)):
 135         ...     print(x)
 136         20..24
 137         18..22
 138         16..28
 139         21..27
 140         """
 141         if self.root is None:
 142             return
 143         yield from self._find_all_overlaps(self.root, x)
 144
 145     def _find_all_overlaps(self, root: bst.Node, x: NumericRange):
 146         if root is None:
 147             return
 148
 149         if root.value.overlaps_with(x):
 150             yield root.value
 151
 152         if root.left:
 153             if root.left.value.highest_in_subtree >= x.low:
 154                 yield from self._find_all_overlaps(root.left, x)
 155
 156         if root.right:
 157             if root.right.value.highest_in_subtree >= x.low:
 158                 yield from self._find_all_overlaps(root.right, x)
 159
 160
 161 if __name__ == "__main__":
 162     import doctest
 163
 164     doctest.testmod()