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})]