[劍指Offer] 54_二叉搜尋樹的第K大節點

題目

給定一棵二叉搜尋樹，請找出其中第k大的節點。

例：

如下圖的二叉搜尋樹，按節點數值大小順序，第三大節點是4。
        5
       / \
      3   7
     / \ / \
    2  4 6  8
           

思路

1. 題目表述有誤，按照其例子應為第K小節點。因為是二叉搜尋樹，是以節點大小是有規律的，以LDR中序序周遊的第K個節點就是第K小節點。
   1. 時間複雜度：O(n)
   2. 空間複雜度：O(logn)

代碼

思路1：時間複雜度：O(n)，空間複雜度：O(logn)

def kth_node_in_BST(root, k):
    """
    :param root: Binary Search Tree root
    :param k: kth smallest
    :return: kth smallest node
    """

    def recursion(root, k):
        nonlocal target
        if not root:
            return k
        k = recursion(root.left, k)
        if not target:
            if k == 1:
                target = root
            else:
                k -= 1
        if not target:
            k = recursion(root.right, k)
        return k
    target = None
    recursion(root, k)
    return target
           

思考

相同的題目

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

題目

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

說明：

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

示例 1:

輸入: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
輸出: 1
           

示例 2:

輸入: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
輸出: 3
           

進階：

如果二叉搜尋樹經常被修改（插入/删除操作）并且你需要頻繁地查找第 k 小的值，你将如何優化 kthSmallest 函數？

代碼

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

class Solution:
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        ans = None
        flag = False
        def search(root,k):
            nonlocal ans
            nonlocal flag
            if flag:
                return 0
            if not root:
                return 0
            count_min = search(root.left, k) + 1
            if  count_min == k:
                flag = True
                ans = root
            r = search(root.right, k - count_min)           
            return count_min + r
        search(root,k)
        return ans.val
           

上面的代碼和正文中不一樣，傳回值為計數，用flag和ans來表征和記錄target。看起來比較繞，沒有後來寫得好。