python3 数据结构和算法(3) 二叉搜索树

import random
import time
from functools import wraps

def timethis(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        elapsed = time.time() - start
        print('%-30s %10s %f' % (func.__name__, ' cost time:', elapsed))
        return result
    return wrapper

class BSTNode():
    def __init__(self, k):
        self.key = k
        self.left = None
        self.right = None
        self.parent = None

    def __str__(self):
        return str(self.key)

class BinarySearchTree():
    def __init__(self):
        self.root = None

    def insert(self, k):
        n = BSTNode(k)
        current = self.root
        parent = None

        while current:
            parent = current
            if k < current.key:
                current = current.left
            else:
                current = current.right

        n.parent = parent
        if not parent:
            self.root = n
        else:
            if k < parent.key:
                parent.left = n
            else:
                parent.right = n

    def inorder_rec(self, node, keys):
        if node:
            self.inorder_rec(node.left, keys)
            keys.append(node.key)
            self.inorder_rec(node.right, keys)

    def inorder(self):
        keys = []
        self.inorder_rec(self.root, keys)
        return keys

    def search(self, k):
        p = self.root
        while p and p.key != k:
            if k < p.key:
                p = p.left
            else:
                p = p.right
        return p

    def minimum(self, node):
        while node.left:
            node = node.left
        return node

    def successor(self, node):
        if node.right:
            return self.minimum(node.right)

        p = node.parent
        while p and node == p.right:
            node = p
            p = p.parent
        return p

    def delete(self, node):
        x = y = None
        if node.left is None or node.right is None:
            y = node
        else:
            y = self.successor(node)

        if y.left is not None:
            x = y.left
        else:
            x = y.right

        if x is not None:
            x.parent = y.parent

        if y.parent is None:
            self.root = x
        else:
            if y == y.parent.left:
                y.parent.left = x
            else:
                y.parent.right = x

        if y != node:
            node.key = y.key


def test_insert(bst):
    for i in range(100):
        k = random.randint(1, 1000)
        if bst.search(k) is None:
            bst.insert(k)

    keys = bst.inorder()
    print(keys)
    s = sorted(keys)
    if s != keys:
        print('error')

def test_search(bst):
    for i in range(10):
        k = random.randint(1, 1000)
        node = bst.search(k)
        if node:
            print(k, 'in tree', node.key)
        else:
            print(k, 'not in tree')

    keys = bst.inorder()
    for i in range(10):
        k = keys[i]
        node = bst.search(k)
        if node:
            print(k, 'in tree', node.key)
        else:
            print(k, 'not in tree')

def test_delete(bst):
    deleted = 0
    while deleted < 10:
        k = random.randint(1, 1000)
        node = bst.search(k)
        if node:
            print('delete node', k)
            bst.delete(node)
            node = bst.search(k)
            if node:
                print('error, should not in tree any more')
                return

            keys = bst.inorder()
            s = sorted(keys)
            if s != keys:
                print('error not sorted after delete node')
            deleted += 1

    # delete all
    keys = bst.inorder()
    for i in range(len(keys)):
        j = random.randint(0, len(keys)-1)
        k = keys[j]
        node = bst.search(k)
        if not node:
            print('error, should in tree')
            return

        print('delete node', k)
        bst.delete(node)
        node = bst.search(k)
        if node:
            print('error, should not in tree any more')
            return

        keys = bst.inorder()
        s = sorted(keys)
        if s != keys:
            print('error not sorted after delete node')


def main():
    bst = BinarySearchTree()
    test_insert(bst)
    test_search(bst)
    test_delete(bst)


if __name__ == '__main__':
    main()