1. 程式人生 > >leet230. 二叉搜尋樹中第K小的元素

leet230. 二叉搜尋樹中第K小的元素

題目:

給定一個二叉搜尋樹,編寫一個函式kthSmallest來查詢其中第k個最小的元素。

注意:
你可以假設k總是有效的,1≤ k ≤二叉搜尋樹元素個數。

進階:
如果經常修改二叉搜尋樹(插入/刪除操作)並且你需要頻繁地找到第k小值呢? 你將如何優化kthSmallest函式?

分析:

  1. 根據BST的特點,任一左子樹節點值不大於根節點值,根節點值不大於右子樹節點值
  2. 採用中序遍歷,節點值依次遞增,當遍歷到第k個節點時剛好為第k小的元素

程式碼:

class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        self.count = 0
        self.ret = 0
        self.flag = True
        def dfs(node,k = k):
            if self.flag and node:
                dfs(node.left)
                self.count += 1
                if self.count == k:
                    self.ret = node.val
                    self.flag = False
                    return self.ret
                dfs(node.right)
        dfs(root,k)
        return self.ret

思考: