BST Mapper

bst_mapper.py

"""Python3 BST Mapper.

Tested on python3.13.7
"""

from __future__ import annotations
from typing import MutableMapping, TypeVar, Iterator, Hashable

K = TypeVar('K', bound=Hashable)
V = TypeVar('V')


class _Node[K, V]:
    __slots__ = ('_key', '_keyh', '_left', '_right', '_value')

    def __init__(self, key: K, value: V) -> None:
        self._key = key
        self._keyh = hash(key)
        self._value = value
        self._left: _Node[K, V] | None = None
        self._right: _Node[K, V] | None = None

    def insert(self, key: K, value: V) -> None:
        keyh = hash(key)
        if keyh == self._keyh:
            self._value = value

        if keyh < self._keyh:
            if self._left:
                self._left.insert(key, value)
            else:
                self._left = _Node(key, value)

        elif keyh > self._keyh:
            if self._right:
                self._right.insert(key, value)
            else:
                self._right = _Node(key, value)

    def search(self, key: K) -> _Node[K, V]:
        keyh = hash(key)
        if self._keyh == keyh:
            return self

        if self._keyh < keyh:
            if not self._right:
                raise KeyError("Key not found")
            return self._right.search(key)

        if not self._left:
            raise KeyError("Key not found")
        return self._left.search(key)

    def iter(self) -> Iterator[_Node[K, V]]:
        if self._left:
            yield from self._left.iter()
        yield self
        if self._right:
            yield from self._right.iter()

    def delete(self, key: K) -> _Node[K, V] | None:
        keyh = hash(key)
        if keyh < self._keyh:
            if self._left:
                self._left = self._left.delete(key)
            else:
                raise KeyError("Key {} not found".format(key))

        elif keyh > self._keyh:
            if self._right:
                self._right = self._right.delete(key)
            else:
                raise KeyError("Key {} not found".format(key))

        elif keyh == self._keyh:
            if not self._left:
                return self._right if self._right else None
            if not self._right:
                return self._left if self._left else None

            successor = self._find_min(self._right)
            self._key = successor._key
            self._keyh = successor._keyh
            self._value = successor._value
            self._right = self._right.delete(successor._key)

        return self

    def _find_min(self, node: _Node[K, V]) -> _Node[K, V]:
        current = node
        while current._left:
            current = current._left
        return current

    def count_nodes(self) -> int:
        count = 1
        if self._left:
            count += self._left.count_nodes()
        if self._right:
            count += self._right.count_nodes()
        return count

    def get_key(self) -> K:
        return self._key

    def get_value(self) -> V:
        return self._value


class BSTMapper(MutableMapping[K, V]):
    def __init__(self) -> None:
        self._root: _Node[K, V] | None = None

    def __getitem__(self, key: K) -> V:
        if not self._root:
            raise KeyError("Key not found")
        return self._root.search(key).get_value()

    def __setitem__(self, key: K, value: V) -> None:
        if not self._root:
            self._root = _Node(key, value)
        else:
            self._root.insert(key, value)

    def __iter__(self) -> Iterator[K]:
        if not self._root:
            return iter([])

        for node in self._root.iter():
            yield node.get_key()

    def __delitem__(self, key: K) -> None:
        if not self._root:
            raise KeyError("Key not found")
        self._root = self._root.delete(key)

    def __len__(self) -> int:
        if not self._root:
            return 0

        return self._root.count_nodes()

    def __str__(self) -> str:
        items: list[str] = []

        for key in self:
            items.append(f"{key}: {self[key]}")

        return "{" + ", ".join(items) + "}"

test.py

import unittest

class TestBSTMapper(unittest.TestCase):
    def test_basic_string_int_mapping(self):
        """Test basic string to int mapping."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2

        self.assertEqual(d['apple'], 1)
        self.assertEqual(d['banana'], 2)
        self.assertEqual(len(d), 2)
        self.assertIn('apple', str(d))
        self.assertIn('banana', str(d))

    def test_mixed_key_types(self):
        """Test mapping with mixed key types (str and int)."""
        d = BSTMapper[str | int, int]()
        d['apple'] = 1
        d['banana'] = 2
        d[5773] = 3

        self.assertEqual(d['apple'], 1)
        self.assertEqual(d['banana'], 2)
        self.assertEqual(d[5773], 3)
        self.assertEqual(len(d), 3)

    def test_nested_bst_mapper(self):
        """Test BSTMapper as a value in another BSTMapper."""
        d = BSTMapper[str | int, BSTMapper[str, int] | str]()
        c = BSTMapper[str, int]()
        c['apple'] = 1
        c['banana'] = 2

        d['node1'] = c
        d['node2'] = c
        d[5773] = "three"

        self.assertEqual(str(d['node1']), str(c))
        self.assertEqual(str(d['node2']), str(c))
        self.assertIn('node1', str(d))
        self.assertIn('node2', str(d))
        self.assertIn('5773', str(d))

    def test_delete_operation(self):
        """Test deletion of keys."""
        d = BSTMapper[str | int, BSTMapper[str, int] | str]()
        c = BSTMapper[str, int]()
        c['apple'] = 1
        c['banana'] = 2

        d['node1'] = c
        d['node2'] = c
        d[5773] = "three"

        del d['node1']

        self.assertEqual(str(d['node2']), str(c))
        self.assertEqual(len(d), 2)
        with self.assertRaises(KeyError):
            _ = d['node1']

    def test_key_not_found(self):
        """Test KeyError is raised for non-existent keys."""
        d = BSTMapper[str, int]()
        with self.assertRaises(KeyError):
            _ = d['nonexistent']

    def test_iteration(self):
        """Test iteration over keys."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2
        d['cherry'] = 3

        keys = list(d)
        self.assertEqual(len(keys), 3)
        self.assertIn('apple', keys)
        self.assertIn('banana', keys)
        self.assertIn('cherry', keys)

    def test_update_existing_key(self):
        """Test updating value for existing key."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['apple'] = 100

        self.assertEqual(d['apple'], 100)
        self.assertEqual(len(d), 1)

    def test_empty_bst_mapper(self):
        """Test empty BSTMapper behavior."""
        d = BSTMapper[str, int]()

        self.assertEqual(len(d), 0)
        self.assertEqual(str(d), "{}")
        self.assertEqual(list(d), [])

    def test_delete_from_empty(self):
        """Test deleting from empty BSTMapper raises KeyError."""
        d = BSTMapper[str, int]()

        with self.assertRaises(KeyError):
            del d['nonexistent']

    def test_delete_only_element(self):
        """Test deleting the only element in BSTMapper."""
        d = BSTMapper[str, int]()
        d['apple'] = 1

        self.assertEqual(len(d), 1)
        del d['apple']

        self.assertEqual(len(d), 0)
        with self.assertRaises(KeyError):
            _ = d['apple']

    def test_delete_nonexistent_key(self):
        """Test deleting non-existent key raises KeyError."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2

        with self.assertRaises(KeyError):
            del d['cherry']

    def test_multiple_deletions(self):
        """Test multiple consecutive deletions."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2
        d['cherry'] = 3
        d['date'] = 4

        self.assertEqual(len(d), 4)

        del d['banana']
        self.assertEqual(len(d), 3)

        del d['cherry']
        self.assertEqual(len(d), 2)

        self.assertEqual(d['apple'], 1)
        self.assertEqual(d['date'], 4)

    def test_overwrite_multiple_times(self):
        """Test overwriting the same key multiple times."""
        d = BSTMapper[str, int]()
        d['key'] = 1
        d['key'] = 2
        d['key'] = 3
        d['key'] = 4

        self.assertEqual(d['key'], 4)
        self.assertEqual(len(d), 1)

    def test_large_dataset(self):
        """Test with a larger dataset."""
        d = BSTMapper[int, int]()
        n = 100

        for i in range(n):
            d[i] = i * 2

        self.assertEqual(len(d), n)

        for i in range(n):
            self.assertEqual(d[i], i * 2)

    def test_delete_root_with_two_children(self):
        """Test deleting root node when it has two children."""
        d = BSTMapper[int, str]()
        d[50] = "fifty"
        d[30] = "thirty"
        d[70] = "seventy"
        d[20] = "twenty"
        d[40] = "forty"
        d[60] = "sixty"
        d[80] = "eighty"

        self.assertEqual(len(d), 7)
        del d[50]
        self.assertEqual(len(d), 6)

        with self.assertRaises(KeyError):
            _ = d[50]

    def test_contains_method(self):
        """Test 'in' operator (uses __contains__ from MutableMapping)."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2

        self.assertIn('apple', d)
        self.assertIn('banana', d)
        self.assertNotIn('cherry', d)

    def test_keys_values_items(self):
        """Test keys(), values(), and items() methods from MutableMapping."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2
        d['cherry'] = 3

        keys = list(d.keys())
        values = list(d.values())
        items = list(d.items())

        self.assertEqual(len(keys), 3)
        self.assertIn('apple', keys)
        self.assertIn(1, values)
        self.assertIn(('banana', 2), items)

    def test_get_method(self):
        """Test get() method with default value from MutableMapping."""
        d = BSTMapper[str, int]()
        d['apple'] = 1

        self.assertEqual(d.get('apple'), 1)
        self.assertEqual(d.get('banana'), None)
        self.assertEqual(d.get('banana', 999), 999)

    def test_pop_method(self):
        """Test pop() method from MutableMapping."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2

        value = d.pop('apple')

        self.assertEqual(value, 1)
        self.assertEqual(len(d), 1)
        self.assertNotIn('apple', d)

    def test_clear_method(self):
        """Test clear() method from MutableMapping."""
        d = BSTMapper[str, int]()
        d['apple'] = 1
        d['banana'] = 2
        d['cherry'] = 3

        d.clear()

        self.assertEqual(len(d), 0)
        self.assertEqual(list(d), [])

unittest.main()