1. 程式人生 > >leetcode 938. Range Sum of BST

leetcode 938. Range Sum of BST

題目

給定一顆二叉搜尋樹的根節點,返回節點值在L和R之間(包括L和R)的節點值的和。

思路

寫了一個完整的二叉搜尋樹,調整一下中序遍歷的方法比較L和R就行了,注意在遞迴呼叫時要寫上L和R這兩個引數。如果是沒有重複數字的BST,可以在遞迴查詢左葉子之後比較一下R和右葉子的值,如果完全相等就早停,速度快了一倍。

程式碼

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None


class Solution:
    def __init__(self):
        self.x= 0

    def rangeSumBST(self, root, L, R):
        """
        :type root: TreeNode
        :type L: int
        :type R: int
        :rtype: int
        """
        if root.left:
            self.rangeSumBST(root.left, L, R)
        if L <= root.val <= R:
            self.x += root.val
            if root.val == R:
                return self.x
        if root.right:
            self.rangeSumBST(root.right, L, R)
        return self.x

其他

二叉搜尋樹的簡單實現

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


class BinarySearchTree:
    def __init__(self):
        self.root = None
        self.elements = []

    def __repr__(self):
        if self.is_empty():
            return 'empty tree'
        else:
            self.elements = []
            self.traverse_mid_first()
            return ' '.join(str(a) for a in self.elements)

    def is_empty(self):
        return False if self.root else True

    def traverse_mid_first(self, node=None, start=float('-inf'), end=float('inf')):
        if not node:
            node = self.root

        if node.left:
            self.traverse_mid_first(node.left, start, end)
        if start <= node.val <= end:
            self.elements.append(node.val)
        if node.right:
            self.traverse_mid_first(node.right, start, end)

    def insert(self, x):
        new_node = TreeNode(x)
        if not self.root:
            self.root = new_node
            return
        current_node, temp_node = self.root, self.root
        while current_node:
            temp_node = current_node
            if x < current_node.val:
                current_node = current_node.left
            else:
                current_node = current_node.right

        if x < temp_node.val:
            temp_node.left = new_node
        if x > temp_node.val:
            temp_node.right = new_node