python 在python中的二进制搜索树

Posted 2021-05-09

# -*- coding: utf-8 -*-
# @Author: cody
# @Date:   2016-11-26 14:18:19
# @Last Modified 2016-11-26
# @Last Modified time: 2016-11-26 17:19:44

from random import shuffle

class Node:
    def __init__(self, val, depth=1):
        self.l = None
        self.r = None
        self.v = val
        self.depth = depth

    def __call__(self):
        """ shorthand way of getting the val """
        return self.v

    def add(self,val,depth=2):
        """ adds an instance to this node """
        if val < self.v: # go to left branch
            if self.l is None: # add if none
                #print("setting {} as left child at depth {}".format(val,depth))
                self.l = Node(val,depth)
            else:
                #print("giving {} to left child({}) at depth {}".format(val,self.l(),depth))
                self.l.add(val,depth+1)
        elif val > self.v: # go to right branch
            if self.r is None:
                #print("setting {} as right child at depth {}".format(val,depth))
                self.r = Node(val,depth)
            else:
                #print("giving {} to right child({}) at depth {}".format(val,self.r(),depth))
                self.r.add(val,depth+1)

    def find(self,target):
        if self() is target:
            return True
        elif self() > target:
            if self.l is None: return False
            else: return self.l.find(target)
        else:
            if self.r is None: return False
            else: return self.r.find(target)



    def map(self):
        """ generator for all nodes """
        yield self
        if self.l is not None:
            for _ in self.l.map():
                yield _
        if self.r is not None:
            for _ in self.r.map():
                yield _

    def __iter__(self):
        return self.map()

    def map_values(self):
        """ generator for all values """
        for _ in self:
            yield _()

    def map_depth(self):
        for _ in self:
            yield _.depth

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

    @property
    def has_root(self):
        return self.root is not None

    def add(self, val):
        if self.has_root:
            self.root.add(val)
        else:
            self.root = Node(val)

    @property
    def shuffled_copy(self):
        """ returns a shuffled copy of itself """
        bt = BinaryTree()
        v=list(self.map_values())
        shuffle(v)
        for _ in v:
            bt.add(_)
        return bt

    @property
    def deepest_branch(self):
        assert self.has_root, "no root found"
        return max(self.root.map_depth())

    def attempt_optimization(self):
        assert self.has_root, "no root found"
        cp = self.shuffled_copy
        if self.deepest_branch > cp.deepest_branch:
            #print("current length: {}".format(self.deepest_branch))
            #print("optimized length: {}".format(cp.deepest_branch))
            self.root = cp.root

    def optimize(self,rounds=100):
        for _ in range(rounds):
            if self.is_optimal:
                break
            self.attempt_optimization()

    @property
    def is_optimal(self):
        """ returns if the tree has been optimized to its fullest. """
        return 2**(self.deepest_branch-1)<self.count

    @property
    def count(self):
        """ returns how many nodes are in the tree """
        out = 0
        for _ in self:
            out += 1
        return out

    def __iter__(self):
        return self.map_values()

    def map_values(self):
        assert self.has_root, "no root found"
        for _ in self.root.map_values():
            yield _

    def find(self,target):
        """ returns True if the instance is there """
        if self.has_root:
            return self.root.find(target)
        else:
            return False


if __name__ == "__main__":
    bt = BinaryTree()
    from random import randint
    r=lambda:randint(1, 100)
    for i in range(50):
        bt.add(r())
    bt.optimize()
    for i in bt:
        print i
    for i in range(50):
        _ = r()
        print _, bt.find(_)