数据结构 | 平衡二叉搜索树 | 问题11

问题描述

给定一个平衡二叉搜索树（AVL树），将其中的每一个节点的键改为该节点原先的键加上所有大于它的键之和。

解决方法

思路

使用中序遍历（inorder traversal）遍历平衡二叉搜索树，对每一个节点进行以下操作：

1. 在平衡二叉搜索树中找到大于该节点键值的所有节点，并将它们的键值相加；
2. 将该节点键值加上步骤1中得到的和。

实现

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

class Solution(object):
    def __init__(self):
        self.sum = 0

    def greater_sum_tree(self, root):
        if root is None:
            return None
        self.greater_sum_tree(root.right)
        root.val += self.sum
        self.sum = root.val
        self.greater_sum_tree(root.left)
        return root

算法复杂度

时间复杂度：O(n)，其中n是平衡二叉搜索树中的节点数，因为每个节点最多被访问一次。

空间复杂度：O(n)，调用递归函数时使用的栈空间。

测试

输入：

root = TreeNode(4)
root.left = TreeNode(1)
root.right = TreeNode(6)
root.left.left = TreeNode(0)
root.left.right = TreeNode(2)
root.right.left = TreeNode(5)
root.right.right = TreeNode(7)
root.left.right.right = TreeNode(3)

输出：

Solution().greater_sum_tree(root)
#     26
#   /   \
# 30     13
# / \   / \
# 30 32 18  7
#        \
#         10

结论

本题使用中序遍历遍历平衡二叉搜索树，并使用一个变量sum记录节点前面的所有节点的和，可以O(n)时间和O(n)空间内完成。