interval_tree.py

from typing import Any, List, Optional

class IntervalNode:
    def __init__(self, low: Any, high: Any, data: Any = None):
        self.low = low
        self.high = high
        self.max = high  # Maximum high in this subtree
        self.data = data
        self.left: Optional[IntervalNode] = None
        self.right: Optional[IntervalNode] = None

class IntervalTree:
    def __init__(self):
        self.root: Optional[IntervalNode] = None

    def insert(self, low: Any, high: Any, data: Any = None) -> None:
        """
        Inserts a new interval [low, high) into the tree.
        """
        new_node = IntervalNode(low, high, data)
        if not self.root:
            self.root = new_node
            return
        
        self._insert_node(self.root, new_node)

    def _insert_node(self, root: IntervalNode, node: IntervalNode) -> None:
        # Standard BST insertion based on low endpoint
        if node.low < root.low:
            if not root.left:
                root.left = node
            else:
                self._insert_node(root.left, node)
        else:
            if not root.right:
                root.right = node
            else:
                self._insert_node(root.right, node)

        # Update the max value of the ancestor node
        if root.max < node.high:
            root.max = node.high

    def overlap_search(self, low: Any, high: Any) -> List[Any]:
        """
        Finds all intervals in the tree that overlap with the query interval [low, high).
        An interval [a, b) overlaps with [x, y) if max(a, x) < min(b, y).
        """
        results = []
        self._overlap_search_node(self.root, low, high, results)
        return results

    def _overlap_search_node(self, root: Optional[IntervalNode], low: Any, high: Any, results: List[Any]) -> None:
        if not root:
            return

        # Check for overlap: max(root.low, low) < min(root.high, high)
        if root.low < high and low < root.high:
            results.append({
                "low": root.low,
                "high": root.high,
                "data": root.data
            })

        # If left child is not empty and its max endpoint is greater than query's low endpoint,
        # then there might be an overlap in the left subtree
        if root.left and root.left.max > low:
            self._overlap_search_node(root.left, low, high, results)

        # We must search the right subtree if root's low endpoint is less than query's high endpoint
        if root.right and root.low < high:
            self._overlap_search_node(root.right, low, high, results)