Leetcode315. 计算右侧小于当前元素的个数 Python实现

• 题目要求：

• 思路：

  • 构建一个二叉树，节点有以下属性：

  • self.left = None
    self.right = None
    self.val = val  # 节点值
    self.nodecount = 1 # 当前节点值出现了多少次
    self.leftcount = 0  # 左子树节点数量

  • 从后往前遍历数组元素，如果当前的数组元素值与当前树节点值相等，那么把当前节点出现的次数加一
  • 如果当前元素的值大于树节点的值，如果当前节点有右节点，继续比较右节点，如果当前节点没有右节点，用当前的数组元素创建新的节点，插入到树中
  • 如果当前元素的值小于节点的值，如果当前节点有左节点，继续比较左节点，如果当前的节点没有左节点，用当前的数组元素创建新的节点，插入到树中
• 代码：

class TreeNode(object):
    def __init__(self, val):
        self.left = None
        self.right = None
        self.val = val  # 节点值
        self.nodecount = 1 # 当前节点值出现了多少次
        self.leftcount = 0  # 左子树节点数量

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        n = len(nums)        
        res = [0] * n
        if n == 0 or n == 1:
            return res

        root = TreeNode(nums[-1])

        for i in range(n - 2, -1, -1):
            res[i] = self.insertNode(root, nums[i])
        return res
    
    # 往二叉搜索树中插入新的节点
    def insertNode(self, root, val):
        cur = root
        res = 0
        while cur:
            if cur.val == val:
                cur.nodecount += 1
                res += cur.leftcount
                break
            elif cur.val < val:
                    res += cur.leftcount + cur.nodecount
                    if cur.right:
                        cur = cur.right
                    else:
                        cur.right = TreeNode(val)
                        break
            else:
                cur.leftcount += 1
                if cur.left:
                    cur = cur.left
                else:
                    cur.left = TreeNode(val)
                    break
        return res

• 完整代码：

class TreeNode(object):
    def __init__(self, val):
        self.left = None
        self.right = None
        self.val = val  
        self.nodecount = 1 
        self.leftcount = 0  

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        n = len(nums)        
        res = [0] * n
        if n == 0 or n == 1:
            return res

        root = TreeNode(nums[-1])

        for i in range(n - 2, -1, -1):
            res[i] = self.insertNode(root, nums[i])
        return res
    
    def insertNode(self, root, val):
        cur = root
        res = 0
        while cur:
            if cur.val == val:
                cur.nodecount += 1
                res += cur.leftcount
                break
            elif cur.val < val:
                    res += cur.leftcount + cur.nodecount
                    if cur.right:
                        cur = cur.right
                    else:
                        cur.right = TreeNode(val)
                        break
            else:
                cur.leftcount += 1
                if cur.left:
                    cur = cur.left
                else:
                    cur.left = TreeNode(val)
                    break
        return res

• 参考视频：LeetCode315 计算右侧小于当前元素的个数 BST解法