(Tree, Medium) Kth Smallest Element in a BST

송재호·2025년 8월 13일
0

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

Binary search tree가 주어졌을 때 k번째로 작은 수를 찾는 문제다.
핵심은 중위 순회다.

Binary search tree 특성 상 중위 순회하면서 count를 증/감 시키면 이것이 몇 번째로 작은 값인지 판단할 수 있게 된다.
(왼쪽 서브트리는 항상 더 작은 값, 오른쪽 서브트리는 항상 더 큰 값)
중위 순회 시 left를 먼저 다 돌고 나온 count가 현재 노드가 몇 번째로 큰 지 알 수 있는 지표

이 문제에서는 k번째라는 명확한 수가 주어졌으므로, cnt = k로 정해놓고 별도의 cnt를 순회마다 하나씩 빼면서 0이 될 때 k번째로 작은 수임을 증명할 수 있다.

class Solution {
    private int answer;
    private int cnt;

    public int kthSmallest(TreeNode root, int k) {
        cnt = k;
        inOrder(root);
        return answer;
    }

    private void inOrder(TreeNode node) {
        if (node == null) {
            return;
        }
        
        // left
        inOrder(node.left);

        // root
        cnt--;
        if (cnt == 0) {
            answer = node.val;
            return;
        }

        // right
        if (cnt > 0) {
            inOrder(node.right);
        }
    }
}

속도가 만족스럽게 안나와서 GPT한테 개선 요청했는데, 스택 기반으로 풀어준다.
아마 느린 재귀 기반이 느린 이유는 메서드 콜 오버헤드, 캐시 지역성, GC 문제??

가독성 면에서 재귀가 더 좋아보이긴 함..

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Deque<TreeNode> stack = new ArrayDeque<>();
        TreeNode cur = root;

        while (cur != null || !stack.isEmpty()) {
            // 1) 왼쪽 끝까지 내려가며 스택에 쌓음
            while (cur != null) {
                stack.push(cur);
                cur = cur.left;
            }
            // 2) 가장 왼쪽 노드를 pop → 방문(작은 값부터 방문됨)
            cur = stack.pop();
            if (--k == 0) {
                return cur.val;
            }
            // 3) 오른쪽 서브트리로 진행
            cur = cur.right;
        }

        // k가 유효하지 않은 경우(문제 조건상 보통 나오지 않음)
        throw new IllegalArgumentException("k is out of bounds");
    }
}
profile
식지 않는 감자

0개의 댓글