python实现二叉查找树

 

class BSTMapNode(object):
    def __init__(self, key, value):
        self.key = key
        self.value = value
        self.left = None
        self.right = None


# 以列表作为底层存储
class BSTMapIterator(object):
    def __init__(self, root):
        self.the_keys = list()
        self.cur_item = 0
        # 初始化内置list，并将索引reset
        self.bst_travel(root)
        self.cur_item = 0

    def __iter__(self):
        return self

    # 遍历列表即可
    def __next__(self):
        if self.cur_item < len(self.the_keys):
            key = self.the_keys[self.cur_item]
            self.cur_item += 1
            return key
        else:
            raise StopIteration

    # 这个递归！中序遍历
    def bst_travel(self, sub_tree):
        if sub_tree is None:
            return

        self.bst_travel(sub_tree.left)
        self.the_keys[self.cur_item] = sub_tree.key
        self.cur_item += 1
        self.bst_travel(sub_tree.right)


class BSTMap(object):
    def __init__(self):
        self._root = None
        self._size = None

    def __len__(self):
        return self._size

    def __iter__(self):
        return BSTMapIterator(self._root)

    def __contains__(self, key):
        return self.bst_search(self._root, key) is not None

    def value_of(self, key):
        node = self.bst_search(self._root, key)
        return node.value

    # 递归搜索
    def bst_search(self, sub_tree, target):
        if sub_tree is None:
            return None
        elif target < sub_tree.key:
            return self.bst_search(sub_tree.left, target)
        elif target > sub_tree.key:
            return self.bst_search(sub_tree.right, target)
        else:
            return sub_tree

    # 查找最小节点
    def bst_min(self, sub_tree):
        if sub_tree is None:
            return None
        elif sub_tree.left is None:
            return sub_tree
        else:
            return self.bst_min(sub_tree.left)

    # 若key=sub_tree的key ,替换value
    def add(self, sub_tree, key, value):
        node = self.bst_search(sub_tree, key)
        if node is not None:
            node.value = value
            return False
        else:
            self._root = self._bst_insert(self._root, key, value)
            self._size += 1
            return True

    # 如果树为空，则将新建节点并返回，若key小于根节点，则递归插入左子树，key大于根节点，递归插入右子树
    def _bst_insert(self, sub_tree, key, value):
        if sub_tree is None:
            sub_tree = BSTMapNode(key, value)
        elif key < sub_tree.key:
            sub_tree.left = self._bst_insert(sub_tree.left, key, value)
        elif key > sub_tree.key:
            sub_tree.right = self._bst_insert(sub_tree.right, key, value)
        return sub_tree

    def remove(self, key):
        self._root = self._bst_remove(self._root, key)
        self._size -= 1

    def _bst_remove(self, sub_tree, target):
        if sub_tree is None:
            return sub_tree
        # 目标比当前节点小，则递归删除左子树的对应节点，否则递归删除右子树对应节点
        elif target < sub_tree.key:
            sub_tree.left = self._bst_remove(sub_tree.left, target)
            return sub_tree
        elif target > sub_tree.key:
            sub_tree.right = self._bst_remove(sub_tree.right, target)
            return sub_tree
        # 目标等于当前节点
        else:
            # 目标是叶子节点
            if sub_tree.left is None and sub_tree.right is None:
                sub_tree = None
            # 目标只有左子树或右子树
            elif sub_tree.left is None or sub_tree.right is None:
                if sub_tree.left is not None:
                    sub_tree = sub_tree.left
                else:
                    sub_tree = sub_tree.right
            # 目标同时有左右子树
            # 查找右子树的最小节点，置为当前节点，并将其删除
            else:
                successor = self.bst_min(sub_tree.right)
                sub_tree.key = successor.key
                sub_tree.value = successor.value
                sub_tree.right = self._bst_remove(sub_tree.right, successor.key)
            return sub_tree