[LeetCode 해석 및 풀이] 230. Kth Smallest Element in a BST(BST의 K번째로 작은 요소)

LeetCode 문제 해석 및 풀이 방법

문제 230. Kth Smallest Element in a BST(BST의 K번째로 작은 요소)

난이도 : 중간🟡 바로가기

 

문제 설명

영문

Given the root of a binary search tree, and an integer k, return the k^th smallest value (1-indexed) of all the values of the nodes in the tree.

 

Follow up: If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?

 

한글

이진 검색 트리의 root와 정수 k 주어지면 트리에 있는 모든 노드 값 중 k^th 로 작은 값 (1-인덱스)을 반환합니다.

 

후속 조치: BST가 자주 수정되고(즉, 삽입 및 삭제 작업을 수행할 수 있음) k번째로 작은 것을 자주 찾아야 한다면 어떻게 최적화하시겠습니까?

 

제한조건

• The number of nodes in the tree is n.
• 1 <= k <= n <= 10⁴
• 0 <= Node.val <= 10⁴

 

입출력 예

입력	출력
root = [3,1,4,null,2], k = 1	1
root = [5,3,6,2,4,null,null,1], k = 3	3

 

해답 및 해설

파이썬 (Python)

첫 번째 풀이

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        class helperOutput:
            def __init__(self, is_answer_in, answer, node_count):
                self.is_answer_in = is_answer_in
                self.answer = answer
                self.node_count = node_count

        def kthHelper(node, k):
            if not node:
                return helperOutput(False, None, 0)

            left_state = kthHelper(node.left, k)
            if left_state.is_answer_in:
                return left_state
            
            present_count = left_state.node_count + 1
            if present_count == k:
                return helperOutput(True, node.val, None)

            right_state = kthHelper(node.right, k-present_count)
            if right_state.is_answer_in:
                return right_state
            
            total_node_count = present_count + right_state.node_count
            return helperOutput(False, None, total_node_count)
        
        answer = kthHelper(root, k).answer

        return answer

왼쪽 트리에서 k번째 수를 찾아본 후, 없으면 k-(왼쪽 트리에 있는 노드의 개수) 번째 수를 찾고, 두 노드의 수의 합이 k보다 작으면 현재 노드 개수를 위로 올려 보내는 방식을 이용하였다. 

 

챗GPT의 풀이

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def in_order_traversal(node):
            if not node:
                return []
            return in_order_traversal(node.left) + [node.val] + in_order_traversal(node.right)
        
        sorted_elements = in_order_traversal(root)
        return sorted_elements[k-1]

 

좀 더 간단한 방법이 없냐고 챗GPT에게 물어보니 다음과 같은 답변을 추천해주었다. BST의 중위 순회(in-order traversal)로 오름차순으로 정렬된 리스트를 만들어 거기서 k번째 수를 찾는 방식이었다. 하지만 이 방법은 간단해 보일지 몰라도 모든 노드를 다 체크해 봐야 되므로 비효율적이다.

 

LeetCode에서 본 풀이

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        result = None
        
        def traverse(root):
            nonlocal result

            if not (root and result is None): 
                return
            
            nonlocal k

            traverse(root.left)
            k -= 1
            if k == 0: 
                result = root.val
            traverse(root.right)

        traverse(root)

        return result

 내 첫번째 풀이에서는 노드의 개수를 세어서 빼주는 것만 생각했는데, 반대로 이 풀이는 노드의 개수를 세면서 하나씩 숫자를 줄여나가는 방식을 이용하였다. 그래서 첫 번째 풀이와 원리는 비슷하지만 더 깔끔한 풀이가 나왔다.