ordered_map.py

from typing import Any, Optional, Iterator, Tuple, Union
from collections.abc import MutableMapping


class MapEntry:
    """Represents a key-value pair in the ordered map."""
    
    def __init__(self, key: Any, value: Any):
        self.key = key
        self.value = value
    
    def __repr__(self) -> str:
        return f"MapEntry({self.key}, {self.value})"
    
    def __eq__(self, other) -> bool:
        if not isinstance(other, MapEntry):
            return False
        return self.key == other.key and self.value == other.value


class OrderedMap(MutableMapping):
    """A sorted dictionary implementation using a binary search tree."""
    
    class _Node:
        """Internal node for the binary search tree."""
        
        def __init__(self, key: Any, value: Any):
            self.key = key
            self.value = value
            self.left: Optional[OrderedMap._Node] = None
            self.right: Optional[OrderedMap._Node] = None
            self.parent: Optional[OrderedMap._Node] = None
    
    def __init__(self):
        self._root: Optional[OrderedMap._Node] = None
        self._size = 0
    
    def __len__(self) -> int:
        return self._size
    
    def __getitem__(self, key: Any) -> Any:
        node = self._find_node(key)
        if node is None:
            raise KeyError(key)
        return node.value
    
    def __setitem__(self, key: Any, value: Any) -> None:
        self.insert(key, value)
    
    def __delitem__(self, key: Any) -> None:
        self.delete(key)
    
    def __iter__(self) -> Iterator[Any]:
        """In-order traversal of keys."""
        for entry in self._inorder_traversal(self._root):
            yield entry.key
    
    def __contains__(self, key: Any) -> bool:
        return self._find_node(key) is not None
    
    def __repr__(self) -> str:
        items = ", ".join(f"{k}: {v}" for k, v in self.items())
        return f"OrderedMap({{{items}}})"
    
    def _find_node(self, key: Any) -> Optional[_Node]:
        """Find node with given key."""
        current = self._root
        while current is not None:
            if key == current.key:
                return current
            elif key < current.key:
                current = current.left
            else:
                current = current.right
        return None
    
    def _find_parent(self, key: Any) -> Tuple[Optional[_Node], bool]:
        """Find parent node and whether key would be left child."""
        if self._root is None:
            return None, False
        
        current = self._root
        parent = None
        is_left = False
        
        while current is not None:
            parent = current
            if key == current.key:
                return parent, True  # Found exact match
            elif key < current.key:
                current = current.left
                is_left = True
            else:
                current = current.right
                is_left = False
        
        return parent, False  # Return parent where key should be inserted
    
    def insert(self, key: Any, value: Any) -> None:
        """Insert a key-value pair into the map."""
        if self._root is None:
            self._root = self._Node(key, value)
            self._size = 1
            return
        
        parent, found = self._find_parent(key)
        if parent is None:
            return  # Should not happen
        
        if found:
            # Key already exists, update value
            parent.value = value
            return
        
        # Create new node
        new_node = self._Node(key, value)
        new_node.parent = parent
        
        # Insert as left or right child
        if key < parent.key:
            parent.left = new_node
        else:
            parent.right = new_node
        
        self._size += 1
    
    def delete(self, key: Any) -> None:
        """Delete a key from the map."""
        node = self._find_node(key)
        if node is None:
            raise KeyError(key)
        
        # Case 1: Node has no children
        if node.left is None and node.right is None:
            if node.parent is None:
                self._root = None
            elif node.parent.left is node:
                node.parent.left = None
            else:
                node.parent.right = None
        
        # Case 2: Node has one child
        elif node.left is None:
            child = node.right
            if node.parent is None:
                self._root = child
            elif node.parent.left is node:
                node.parent.left = child
            else:
                node.parent.right = child
            if child:
                child.parent = node.parent
        
        elif node.right is None:
            child = node.left
            if node.parent is None:
                self._root = child
            elif node.parent.left is node:
                node.parent.left = child
            else:
                node.parent.right = child
            if child:
                child.parent = node.parent
        
        # Case 3: Node has two children
        else:
            # Find successor (smallest key in right subtree)
            successor = node.right
            while successor.left is not None:
                successor = successor.left
            
            # Replace node's key and value with successor's
            node.key = successor.key
            node.value = successor.value
            
            # Delete successor
            if successor.parent.left is successor:
                successor.parent.left = successor.right
            else:
                successor.parent.right = successor.right
            if successor.right:
                successor.right.parent = successor.parent
        
        self._size -= 1
    
    def _inorder_traversal(self, node: Optional[_Node]) -> Iterator[MapEntry]:
        """In-order traversal generator."""
        if node is not None:
            yield from self._inorder_traversal(node.left)
            yield MapEntry(node.key, node.value)
            yield from self._inorder_traversal(node.right)
    
    def lower_bound(self, key: Any) -> Optional[MapEntry]:
        """Find the smallest entry with key >= given key."""
        result = None
        current = self._root
        
        while current is not None:
            if current.key >= key:
                result = MapEntry(current.key, current.value)
                current = current.left
            else:
                current = current.right
        
        return result
    
    def upper_bound(self, key: Any) -> Optional[MapEntry]:
        """Find the smallest entry with key > given key."""
        result = None
        current = self._root
        
        while current is not None:
            if current.key > key:
                result = MapEntry(current.key, current.value)
                current = current.left
            else:
                current = current.right
        
        return result
    
    def range(self, start_key: Any, end_key: Any) -> Iterator[MapEntry]:
        """Iterate over entries with keys in [start_key, end_key)."""
        current = self.lower_bound(start_key)
        if current is None:
            return
        
        # Find the starting node
        node = self._find_node(current.key)
        if node is None:
            return
        
        # In-order traversal starting from node, stopping at end_key
        for entry in self._inorder_traversal_from(node):
            if entry.key >= end_key:
                break
            yield entry
    
    def _inorder_traversal_from(self, node: _Node) -> Iterator[MapEntry]:
        """In-order traversal starting from a specific node."""
        # This is a simplified version that just does full traversal
        # A more efficient implementation would start from the given node
        yield from self._inorder_traversal(self._root)


def main():
    """Demo the OrderedMap functionality."""
    print("Creating OrderedMap...")
    omap = OrderedMap()
    
    # Insert some data
    data = [(5, "five"), (2, "two"), (8, "eight"), (1, "one"), (3, "three"), 
            (7, "seven"), (9, "nine"), (4, "four"), (6, "six")]
    
    print("Inserting data:")
    for key, value in data:
        omap[key] = value
        print(f"  Inserted {key}: {value}")
    
    print(f"\nMap size: {len(omap)}")
    print(f"Map contents: {omap}")
    
    # Test access
    print(f"\nAccessing key 5: {omap[5]}")
    
    # Test lower_bound and upper_bound
    print(f"\nLower bound of 4: {omap.lower_bound(4)}")
    print(f"Upper bound of 4: {omap.upper_bound(4)}")
    print(f"Lower bound of 10: {omap.lower_bound(10)}")
    print(f"Upper bound of 0: {omap.upper_bound(0)}")
    
    # Test range iteration
    print("\nRange [3, 7):")
    for entry in omap.range(3, 7):
        print(f"  {entry}")
    
    # Test deletion
    print(f"\nDeleting key 5...")
    del omap[5]
    print(f"Map after deletion: {omap}")
    print(f"Map size: {len(omap)}")
    
    # Test contains
    print(f"\nContains key 5: {5 in omap}")
    print(f"Contains key 2: {2 in omap}")
    
    # Test iteration
    print("\nIterating through keys:")
    for key in omap:
        print(f"  {key}: {omap[key]}")


if __name__ == "__main__":
    main()