LeetCode 13 二叉搜索樹中第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小的元素，即現將所有元素用中序遍歷排序後輸出下標第k-1個元素
  後來又發現；列表是沒有必要的

代碼

#用列表中轉中序遍歷
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    node_list=[]
    def inorder(self,head):
        if head==None:
            return
        self.inorder(head.left)
        self.node_list.append(head.val)
        self.inorder(head.right)
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.node_list=[]
        self.inorder(root)
        return self.node_list[k-1]
        

#改進版本
class Solution(object):
    def kthSmallest(self, root, k):
        self.count=k
        self.res=0
        def core(root):
            if  root :
                core(root.left)
                self.count=self.count-1
                if self.count==0:
                    self.res=root.val 
                core(root.right)

        core(root)
        return self.res