[leetcode] 230. Kth Smallest Element in a BST 找出二叉搜索樹中的第k小的元素

題目大意

https://leetcode.com/problems/kth-smallest-element-in-a-bst/description/

230. Kth Smallest Element in a BST

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note:
You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Example 1:

Input: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
Output: 3

Follow up:

What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

 

給定一棵二叉搜索樹（BST），編寫一個函數kthSmallest找出其中第k小的元素。

注意：
你能夠假設k老是有效的， 1 ≤ k ≤ BST的元素總數。

進一步思考：
若是BST的修改（插入/刪除）操做十分頻繁，而且須要頻繁地找出第k小的元素，應該怎樣優化kthSmallest函數？

 

解題思路

BST具備以下性質：

• 左子樹中全部元素的值均小於根節點的值
• 右子樹中全部元素的值均大於根節點的值

所以採用中序遍歷（左 -> 根 -> 右）便可以遞增順序訪問BST中的節點，從而獲得第k小的元素，時間複雜度O(k)

Python代碼：

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

class Solution(object):      
    def kthSmallest(self, root, k):  # 52 ms
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        stack = []
        node = root
        while node:
            stack.append(node)
            node = node.left
        x = 1
        while stack and x <= k:
            node = stack.pop()
            x += 1
            right = node.right
            while right:
                stack.append(right)
                right = right.left
        return node.val

 

遞歸方式：

class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        cnt = []
        self.helper(root, cnt, k)
        return cnt[k - 1]

    def helper(self, node, cnt, k):
        if not node:
            return None
        self.helper(node.left, cnt, k)
        cnt.append(node.val)
        if len(cnt) == k:  # 56 ms  <= 96ms
            return None
        self.helper(node.right, cnt, k)

 

進一步思考：
若是BST節點TreeNode的屬性能夠擴展，則再添加一個屬性leftCnt，記錄左子樹的節點個數

記當前節點爲node
當node不爲空時循環：
若k == node.leftCnt + 1：則返回node
不然，若k > node.leftCnt：則令k -= node.leftCnt + 1，令node = node.right
不然，node = node.left

上述算法時間複雜度爲O(BST的高度)

 

參考

http://bookshadow.com/weblog/2015/07/02/leetcode-kth-smallest-element-bst/

https://leetcode.com/problems/kth-smallest-element-in-a-bst/discuss/63660/3-ways-implemented-in-JAVA-(Python):-Binary-Search-in-order-iterative-and-recursive