蓝桥杯：BST插入节点问题之Python求解

【问题描述】
给定一棵包含 N 个节点的二叉树，节点编号是 1 ∼ N。其中 i 号节点具有权值 W i ，并且这些节点的权值恰好形成了一棵排序二叉树 (BST)。
现在给定一个节点编号 K，小明想知道，在这 N 个权值以外，有多少个整数 X (即 X 不等于任何 W i ) 满足：给编号为 K 的节点增加一个权值为 X 的子节点，仍可以得到一棵 BST。
例如在下图中，括号外的数字表示编号、括号内的数字表示权值。即编号1 ∼ 4 的节点权值依次是 0、10、20、30。

如果 K = 1，那么答案为 0。因为 1 号节点已经有左右子节点，不能再增加子节点了。
如果 K = 2，那么答案为无穷多。因为任何一个负数都可以作为 2 的左子节点。
如果 K = 3，那么答案为 9。因为 X = 11,12,··· ,19 都可以作为 3 的左子节点。

【输入格式】
第一行包含 2 个整数 N 和 K。
以下 N 行每行包含 2 个整数，其中第 i 行是编号为 i 的节点的父节点编号P i 和权值 W i 。注意 P i = 0 表示 i 是根节点。
输入保证是一棵 BST。

【输出格式】
一个整数代表答案。如果答案是无穷多，输出 −1。

【样例输入】
4 3
0 10
1 0
1 20
3 30
【样例输出】
9

【评测用例规模与约定】
对于 60% 的评测用例，1 ≤ K ≤ N ≤ 100，0 ≤ W i ≤ 200，且 W i 各不相同。
对于所有评测用例，1 ≤ K ≤ N ≤ 10000，0 ≤ W i ≤ 100000000，且 W i 各不相同。

我的解题思路：
这是一道关于二叉树的题，如果是用C/C++，可以考虑用链表求解，但在Python中没有链表，只能考虑用列表、字典等数据结构。
要为节点K新增一个权值为X的子节点，使新生成的二叉树仍然是BST二叉树，求X可能的值有多少种（用result来表示），可能碰到如下几种情况：
1.节点K左右子树都存在，则result=0
2.节点K有左子树无右子树，新插入的节点充当节点K的右子树，此时又分两种情况：
（1）节点K为其父节点的左子树，那么有W(k) （2）节点K为其父节点的右子树，那么有W§W(k)即可：result=-1
3.节点K有右子树无左子树，新插入的节点充当节点K的左子树，此时又分两种情况：
（1）节点K为其父节点的左子树，那么有X （2）节点K为其父节点的右子树，那么有W§ 4.节点K既无左子树也无右子树，则result=新节点为左子树时的X取值情况+新节点为右子树时的X取值情况，不管K是其父节点的左子树还是右子树，result的两个加数总有一个趋近正无穷，因而result也趋近正无穷，故result=-1

以上四种情况，不管是在哪种情况下，要求解result，都需判断K节点的子树情况，在情况2和情况3下，还需要判断K节点是属于右子树还是左子树，这解需要知道K节点父节点的权值。此外，这两种情况还需判断K节点是不是根节点。

因此，我用了三个字典分别存节点的父节点信息、权值、以及子树信息，再用两个函数来分别创建树和插入子树，具体代码如下：

DParent=dict()   #存节点的父节点信息，{节点编号：父节点}
DWeight=dict()   #存节点的权重信息，{节点编号：权重}
DChildren=dict() #存节点的子树信息，{节点编号：[左子节点，右子节点]}

Str=input()
tempL=Str.split(' ')
N=int(tempL[0])
K=int(tempL[1])

def build_tree():
    for i in range(1,N+1):
        Str=input()
        tempL=Str.split(' ')
        P=int(tempL[0])
        W=int(tempL[1])
        DParent[i]=P
        DWeight[i]=W
        DChildren[i]=[-1,-1]
        
        '''比较节点 i 与其父节点的权值，以判断它是左还是右子树，并将此信息写入其父节点的子树字典中'''
        if i>1:   
            if W>DWeight[P]:
                DChildren[P][1]=i  #节点 i 为其父节点的右子树
            if W<DWeight[P]:
                DChildren[P][0]=i  #节点 i 为其父节点的左子树
            
def insert_subtree():
    LChild=DChildren[K][0]   #节点K的左子树信息
    RChild=DChildren[K][1]   #节点K的右子树信息
    #1.节点K既无左子树也无右子树
    if LChild>-1 and RChild>-1:
        result=0
    #2.节点K有左子树无右子树
    elif LChild>-1 and RChild==-1:  #插入点为右子树
        P=DParent[K]
        if P==0:         
            result=-1
        else:    
            if DWeight[K]>DWeight[P]:  #K节点属于右子树
                result=-1
            else:                      #K节点属于左子树
                result=DWeight[P]-DWeight[K]-1
    #3.节点K无左子树有右子树
    elif LChild==-1 and RChild>-1:  #插入点为左子树
        P=DParent[K]
        if P==0:
            result=-1
        else:
            if DWeight[K]>DWeight[P]:  #K节点属于右子树
                result=DWeight[K]-DWeight[P]-1
            else:                      #K节点属于左子树
                result=-1
    #4.节点K既无左子树也无右子树
    elif LChild==-1 and RChild==-1:  
        result=-1
    print(result)
    

def main():
    build_tree()
    insert_subtree()

main()