python实现二分搜索树以及改进后的AVL Tree

二分搜索树讲解
python代码:

class TreeNode(object):
    def __init__(self, key, val, parent=None, left_child=None, right_child=None):
        self.key = key 
        self.val = val 
        self.parent = parent # 指向该节点的双亲节点
        self.left_child = left_child 
        self.right_child = right_child

    def is_left_child(self):
        return self.parent and self.parent.left_child == self

    def is_right_child(self):
        return self.parent and self.parent.right_child == self

    def is_root(self):
        return not self.parent

    def is_leaf(self):
        return not (self.left_child or self.right_child)

    def has_any_child(self):
        return self.left_child or self.right_child

    def has_both_childs(self):
        return self.left_child and self.right_child

    def find_successor(self):
        """
        寻找该节点的successor
        successor必须是leaf node 或 只有一个子节点
        """
        if self.right_child:
            return self.right_child.find_min()
        elif self.is_left_child():
            return self.parent
        else:
            self.parent.right_child = None
            succ = self.parent.find_successor()
            self.parent.right_child = self
            return succ

    def find_min(self):
        # 寻找该节点及其子节点中key最小的节点
        current_node = self
        while current_node.left_child:
            current_node = current_node.left_child
        return current_node

    def splice(self):
        """
        去除successor,并将其子节点和其双亲节点嫁接
        """
        if self.is_leaf():
            if self.is_left_child():
                self.parent.left_child = None
            else:
                self.parent.right_child = None
        elif self.has_any_child():
            if self.left_child:
                if self.is_left_child():
                    self.parent.left_child = self.left_child
                else:
                    self.parent.right_child = self.left_child
                self.left_child.parent = self.parent
            else:
                if self.is_left_child():
                    self.parent.left_child = self.right_child
                else:
                    self.parent.right_child = self.right_child
                self.right_child.parent = self.parent

    def replace_node_data_childs(self, new_key, new_val, new_left_child, new_right_child):
        self.key = new_key
        self.val = new_val
        self.left_child = new_left_child
        self.right_child = new_right_child

    def __iter__(self):
        if self:
            if self.left_child:
                for elem in self.left_child:
                    yield elem
            yield self.key
            if self.right_child:
                for elem in self.right_child:
                    yield elem


class BinSearchTree(object):
    def __init__(self, root=None):
        self.root = None
        self._size = 0

    def __len__(self):
        return self._size

    def __setitem__(self, k, v):
        if self.root:
            self.__put(k, v, self.root)
        else:  # the search tree is empty
            self.root = TreeNode(k, v)
        self._size += 1

    def __put(self, k, v, current_node):
        if k < current_node.key:
            if current_node.left_child:
                self.__put(k, v, current_node.left_child)
            else:
                current_node.left_child = TreeNode(k, v, parent=current_node)
        elif k > current_node.key:
            if current_node.right_child:
                self.__put(k, v, current_node.right_child)
            else:
                current_node.right_child = TreeNode(k, v, parent=current_node)
        else:  # the k == current_node.key
            current_node.val = v

    def __getitem__(self, k):
        if self._size > 1:
            item = self.__get(k, self.root)
            if item:
                return self.__get(k, self.root).val
            else:
                return None
        elif self._size == 1 and self.root.key == k:
            return self.root.val
        else:
            return None

            # if self.root:
            #   if self._size == 1: # the tree only has root node
            #       if self.root.key == k:
            #           return self.root.val
            #       else:
            #           return None
            #   else:
            #       if self.__get(k, self.root):
            #           return self.__get(k, self.root).val
            #       return None
            # else: # the tree is empty
            #   return None

    def __get(self, k, current_node):
        if not current_node:
            return None
        elif k == current_node.key:
            return current_node
        elif k < current_node.key:
            return self.__get(k, current_node.left_child)
        else:
            return self.__get(k, current_node.right_child)

            # if current_node:
            #   if k < current_node.key:
            #       return self.__get(k, current_node.left_child)
            #   elif k > current_node.key:
            #       return self.__get(k, current_node.right_child)
            #   else:
            #       return current_node
            # else: # the tree doesn't has the key
            #   return None

    def __contains__(self, k):
        if self.__getitem__(k):
            return True
        else:
            return False

    def __delitem__(self, k):
        if self._size > 1:
            del_node = self.__get(k, self.root)
            if del_node:
                self.__remove(del_node)
                self._size -= 1
            else:
                raise KeyError('Error, key not in tree')
        elif self._size == 1 and self.root.val == k:
            self.root = None
            self._size -= 1
        else:
            raise KeyError('Error, key not in tree')

    def __remove(self, del_node):
        # if the node that would be deleted has no childs
        if del_node.is_leaf():
            if del_node.is_left_child():
                del_node.parent.left_child = None
            else:
                del_node.parent.right_child = None
        elif del_node.has_both_childs():
            successor = del_node.find_successor()
            successor.splice()
            del_node.key = successor.key
            del_node.val = successor.val
        else:  # the node that would be deleted has a child
            if del_node.left_child:
                if del_node.is_left_child():
                    del_node.parent.left_child, del_node.left_child.parent = \
                        del_node.left_child, del_node.parent
                elif del_node.is_right_child():
                    del_node.parent.right_child, del_node.left_child.parent = \
                        del_node.left_child, del_node.parent
                else:  # the node that would be deleted is root node
                    del_node.replace_node_data_childs(del_node.left_child.key,
                                                      del_node.left_child.val,
                                                      del_node.left_child.left_child,
                                                      del_node.left_child.right_child)

            else:  # the node that would be deleted has right child
                if del_node.is_left_child():
                    del_node.parent.left_child, del_node.right_child.parent = \
                        del_node.right_child, del_node.parent
                elif del_node.is_right_child():
                    del_node.parent.right_child, del_node.right_child.parent = \
                        del_node.right_child, del_node.parent
                else:
                    del_node.replace_node_data_childs(del_node.right_child.key,
                                                      del_node.right_child.val,
                                                      del_node.right_child.left_child,
                                                      del_node.right_child.right_child)

    def __iter__(self):
        if self.root:
            for k in self.root:
                yield k
        else:
            raise KeyError('the tree is empty')


if __name__ == '__main__':
    mytree = BinSearchTree()
    # setitem
    mytree[3] = "red"
    mytree[4] = "blue"
    mytree[6] = "yellow"
    mytree[2] = "at"
    mytree[5] = "green"
    # delete item
    del mytree[4]

    for k in mytree:
        print(k)

AVL树讲解
python代码:

from bin_search_tree import BinSearchTree, TreeNode

class AVLTreeNode(TreeNode):
    def __init__(self, key, val, parent=None, left_child=None, right_child=None):
        super().__init__(key, val, parent, left_child, right_child)
        self.balance_factor = 0

class AVLTree(BinSearchTree):
    def __init__(self, root=None):
        super().__init__(root)

    def __setitem__(self, k, v):
        """
        override the super().__setitem__
        """
        if self.root:
            self.__put(k, v, self.root)
        else:  # the search tree is empty
            self.root = AVLTreeNode(k, v)
        self._size += 1

    def __put(self, k, v, current_node):
        """
        override the super().__put
        """
        if k < current_node.key:
            if current_node.left_child:
                self.__put(k, v, current_node.left_child)
            else:
                current_node.left_child = AVLTreeNode(k, v, parent=current_node)
                self.__update_balance(current_node.left_child)
        elif k > current_node.key:
            if current_node.right_child:
                self.__put(k, v, current_node.right_child)
            else:
                current_node.right_child = AVLTreeNode(k, v, parent=current_node)
                self.__update_balance(current_node.right_child)
        else:
            current_node.val = v

    def __update_balance(self, current_node):
        if current_node.balance_factor < -1 or \
           current_node.balance_factor > 1:
           self.__rebalance(current_node)
           return

        if current_node.parent != None:
            if current_node.is_left_child():
                current_node.parent.balance_factor += 1
            elif current_node.is_right_child():
                current_node.parent.balance_factor -= 1

            if current_node.parent.balance_factor != 0:
                self.__update_balance(current_node.parent)

    def __rebalance(self, current_node):
        if current_node.balance_factor < -1: # need left rotation
            if current_node.right_child.balance_factor > 0:
                    self.__right_rotate(current_node.right_child)
            self.__left_rotate(current_node)
        elif current_node.balance_factor > 1: # need right rotation
            if current_node.left_child.balance_factor < 0:
                self.__left_rotate(current_node.left_child)
            self.__right_rotate(current_node)

    def __left_rotate(self, rot_root):
        """
        the node's balance factor < -1
        """
        # 将rot root的右子节点设为new root
        new_root = rot_root.right_child

        # 将new root的左子节点设置为rot root的右子节点
        rot_root.right_child = new_root.left_child
        if new_root.left_child:
            new_root.left_child.parent = rot_root

        # 将new root的双亲节点设置为rot root的双亲节点
        new_root.parent = rot_root.parent
        if rot_root.is_root():
            self.root = new_root
        else:
            if rot_root.is_left_child():
                rot_root.parent.left_child = new_root
            else:
                rot_root.parent.right_child = new_root

        # 将rot root的双亲节点设置为new root(rot root为new root的左子节点)
        new_root.left_child = rot_root
        rot_root.parent = new_root

        # critical steps
        # update the balance factor both of rot root and new root
        rot_root.balance_factor = rot_root.balance_factor + 1 - min(new_root.balance_factor, 0)
        new_root.balance_factor = new_root.balance_factor + 1 + max(rot_root.balance_factor, 0)

    def __right_rotate(self, rot_root):
        """
        the node's balance factor > 1
        """
        # 将rot root的左子节点设为new root
        new_root = rot_root.left_child

        # 将new root的右子节点设置为rot root的左子节点
        rot_root.left_child = new_root.right_child
        if new_root.right_child:
            new_root.right_child.parent = rot_root

        # 将new root的双亲节点设置为rot root的双亲节点
        new_root.parent = rot_root.parent
        if rot_root.is_root():
            self.root = new_root
        else:
            if rot_root.is_left_child():
                rot_root.parent.left_child = new_root
            else:
                rot_root.parent.right_child = new_root

        # 将rot root的双亲节点设置为new root(rot root为new root的右子节点)
        new_root.right_child = rot_root
        rot_root.parent = new_root

        # critical steps
        # update the balance factor both of rot root and new root
        rot_root.balance_factor = rot_root.balance_factor - 1 + min(new_root.balance_factor, 0)
        new_root.balance_factor = new_root.balance_factor - 1 - max(rot_root.balance_factor, 0)


if __name__ == '__main__':
    mytree = AVLTree()
    # setitem
    mytree[3] = "red"
    mytree[4] = "blue"
    mytree[6] = "yellow"
    # mytree[2] = "at"
    # mytree[5] = "green"

    print('root:%s|balance_factor:%s' % (mytree.root.key, mytree.root.balance_factor))
    print('root.left_child:%s' % mytree.root.left_child.key)
    print('root.right_child:%s' % mytree.root.right_child.key)
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值