BST二叉搜索树插入节点建树并找出不平衡节点，networkx，Python

梦里梦外; 2023-10-02 16:30 1阅读 0赞

BST二叉搜索树插入节点建树并找出失衡节点，networkx，Python

import random
    
    from matplotlib import pyplot as plt
    import networkx as nx
    
    PARENT = 'parent'
    LEFT = 'left'
    RIGHT = 'right'
    
    
    def app():
        SIZE = 5
        data = []
        for i in range(SIZE):
            data.append(i)
        random.shuffle(data)
    
        G = nx.DiGraph()
        while len(data) > 0:
            print('-')
            d = data.pop()
            insert(G, d)
    
            unblance_node = check_blance(G)
            if unblance_node is not None:
                print('找到失衡点', unblance_node)
    
        print(G.nodes(data=True))
    
        pos = nx.spring_layout(G)
        nx.draw(G, pos,
                node_color='red',
                node_size=300,
                font_size=10,
                font_color='green',
                with_labels=True)
    
        plt.show()
    
    
    def check_blance(G):
        node = None
        for n in G.nodes(data=True):
            f = get_blance_factor(G, n[0])
            if abs(f) > 1:
                node = (n, f)
        return node
    
    
    def get_blance_factor(G, number):
        factor = 0
        if get_node_height(G, number) == 0:
            factor = 0
            # print(number, '平衡因子', factor)
            return factor
    
        left = G.nodes[number][LEFT]
        right = G.nodes[number][RIGHT]
    
        lh = 0
        rh = 0
        if left is not None:
            lh = get_node_height(G, left) + 1
    
        if right is not None:
            rh = get_node_height(G, right) + 1
    
        factor = lh - rh
        # print(number, '平衡因子', factor)
        return factor
    
    
    def get_node_height(G, number):
        height = 0
        node = get_node_by_value(G, number)
        if node[1][LEFT] is None and node[1][RIGHT] is None:
            # print(number, '高度', height)
            return height
    
        bfs = nx.bfs_tree(G, source=number)
        last = list(bfs).pop()
        while True:
            last_n = get_node_by_value(G, last)
            if last_n[1][PARENT] is None:
                break
            else:
                height = height + 1
            last = last_n[1][PARENT]
            if last == number:
                break
        # print(number, '高度', height)
        return height
    
    
    def get_node_by_value(G, v):
        node = None
        for n in G.nodes(data=True):
            if n[0] == v:
                node = n
                break
        return node
    
    
    def insert(G, d):
        if G.number_of_nodes() == 0:
            G.add_node(d)
            G.nodes[d][PARENT] = None
            G.nodes[d][LEFT] = None
            G.nodes[d][RIGHT] = None
            return
    
        print('开始插入', d)
        root_node = None
        for n in G.nodes(data=True):
            if n[1][PARENT] is None:
                root_node = n
                break
        print('根节点', root_node)
    
        while True:
            left = root_node[1][LEFT]
            right = root_node[1][RIGHT]
    
            if left is None:
                if d < root_node[0]:
                    root_node[1][LEFT] = d
                    G.add_node(d)
                    G.nodes[d][PARENT] = root_node[0]
                    G.nodes[d][LEFT] = None
                    G.nodes[d][RIGHT] = None
    
                    G.add_edge(root_node[0], d)
                    break
    
            if right is None:
                if d > root_node[0]:
                    root_node[1][RIGHT] = d
                    G.add_node(d)
                    G.nodes[d][PARENT] = root_node[0]
                    G.nodes[d][LEFT] = None
                    G.nodes[d][RIGHT] = None
    
                    G.add_edge(root_node[0], d)
                    break
    
            if d > root_node[0]:
                val = root_node[1][RIGHT]
                root_node = get_node_by_value(G, val)
            else:
                val = root_node[1][LEFT]
                root_node = get_node_by_value(G, val)
    
    
    if __name__ == '__main__':
        app()

输出：

-
    -
    开始插入 4
    根节点 (3, {'parent': None, 'left': None, 'right': None})
    -
    开始插入 2
    根节点 (3, {'parent': None, 'left': None, 'right': 4})
    -
    开始插入 0
    根节点 (3, {'parent': None, 'left': 2, 'right': 4})
    -
    开始插入 1
    根节点 (3, {'parent': None, 'left': 2, 'right': 4})
    找到失衡点 ((2, {'parent': 3, 'left': 0, 'right': None}), 2)
    [(3, {'parent': None, 'left': 2, 'right': 4}), (4, {'parent': 3, 'left': None, 'right': None}), (2, {'parent': 3, 'left': 0, 'right': None}), (0, {'parent': 2, 'left': None, 'right': 1}), (1, {'parent': 0, 'left': None, 'right': None})]