Scribbling

LeetCode: 230. Kth Smallest Element in a BST 본문

Computer Science/Coding Test

LeetCode: 230. Kth Smallest Element in a BST

focalpoint 2021. 12. 14. 21:33

BST이기 때문에 Inorder-Traverse하면 된다.

K번째 원소를 찾으면 순회를 중단한다. 때문에 O(K)이다.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        self.ret = 0
        self.k = k
        self.traverse(root)
        return self.ret
    
    # returns whether traverse should be continued
    def traverse(self, node):
        if node.left:
            if self.traverse(node.left):
                return True
        self.k -= 1
        if self.k == 0:
            self.ret = node.val
        if node.right:
            if self.traverse(node.right):
                return True
        return False