230. Kth Smallest Element in a BST

정보

• 문제 보기: 230. Kth Smallest Element in a BST
• 소요 시간: 17분 29초
• 풀이 언어: python
• 체감 난이도: 2️⃣
• 리뷰 횟수: ✅

풀이 키워드

스포주의

dfs

풀이 코드

정보

• 메모리: 20940 KB
• 시간: 0 ms

# 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

N, K = 0, 0
ans = 0

def rcs(root):
    global N, ans
    # base case
    if root is None:
        return

    rcs(root.left)
    N += 1
    if N == K:
        ans = root.val
        raise
    rcs(root.right)

    return

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        global N, K, ans
        N = 0
        K = k
        ans = 0
        try:
            rcs(root)
        except:
            pass
        return ans

풀이 해설

BST를 중위 순회하여 k번째 녀석을 반환하면 되는 문제이다.

base case는 도착지가 null일 때이고, 이는 다음의 condition으로 판단했다.

# base case
if root is None:
    return

팁

파이썬에서 즉시 재귀 탈출하기

파이썬은 goto나 label이 없어 즉시 탈출이 어렵다.

하지만 딱 하나 방법이 있었으니...

if N == K:
    ans = root.val
    raise

try:
    rcs(root)
except:
    pass

처음엔 rcs가 재귀함수이기 때문에, 정답을 찾아서 return해도 스택 정리하느라 4ms가 걸렸다.

따라서 예외 처리 트릭을 이용하여 즉시 재귀를 탈출하도록 수정했고, 수행시간이 0ms로 최적화되었다.

메모

• 리트코드 채점 방식이 while (tc--) 와 같은 방식이다.
  • 솔루션 함수 내에서 N과 ans를 초기화해주어야 함.