[C++/LeetCode] Kth Smallest Element in a BST

Daily LeetCoding Challenge for April, Day 18.

 

문제

Binary Search Tree에서 K번째로 작은 수를 구하는 것이다.

 

아이디어1: min heap

binary sarch tree의 아이템을 순회하여 다 가져와서 min heap(priority queue)에 넣는다.그 후 priority queue에서 K번째의 아이템을 가져오면 됨

 

완성

submission: 21ms / 24.5MB

 

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    void fillPQ(TreeNode *root, priority_queue<int, vector<int>, greater<int>> &pq) {
        if (root == NULL)
            return;
        pq.push(root->val);
        fillPQ(root->left, pq);
        fillPQ(root->right, pq);
    }
    
    int kthSmallest(TreeNode* root, int k) {
        priority_queue<int, vector<int>, greater<int>> pq;
        
        fillPQ(root, pq);

        for(int i = 0;i < k - 1;i++)
            pq.pop();
        
        return pq.top();
    }
};

 

이 방법은 단순하고 이해하기 쉽지만 큰 문제가 있다.

바로 다음과 같은 경우일 때이다.

BST가 자주 수정된다면 이를 private 멤버 변수로 선언할 수도 없고 (자주 변경되므로 효율면에서 굳이 저장할 필요가 없음)

kth 작은 수를 자주 호출한다면 매번 priority queue 자료구조를 생성하는 것도 문제다.

게다가 트리를 항상 전부 확인해야 하기 때문에 트리 구조가 커질수록 시간도 오래 걸리고 메모리 효율도 별로다.

따라서 다른 아이디어가 필요하다.

 

 

아이디어2: in-order traversal

문제에서는 Kth 작은 수를 찾는 것이 목표이다.

즉 root->left 에서 수를 찾고, 없으면 root 확인, 그래도 없으면 root->right 에서 찾는 것이다.

즉 tree를 in-order traversal를 하여 찾을 수 있다.

이를 간단하게 재귀를 통해서 구현할 수 있다.

 

완성

submission: 16ms / 24.1MB

 

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
private:
    int index = 0;
    int result;
public:
    void doSearch(TreeNode *root, int k) {
        if (root == NULL)
            return;
        doSearch(root->left, k);
        index++;
        if (index == k) {
            result = root->val;
            return;
        }
        doSearch(root->right, k);
    }
    int kthSmallest(TreeNode* root, int k) {
        doSearch(root, k);
        return result;
    }
};

 

private에 선언된 index는 현재 몇번째 수를 확인했는지 알려주는 역할을 하며, result 는 결과값이다.

 

먼저 root == NULL 인 경우에는 동작하지 않음을 확인하자.

그 다음 먼저 root->left 에서 Kth 수를 찾고 있다. 이때 index를 증가시키면서 몇번째로 작은 수를 확인한 것인지 체크한다.

체크한 index가 k와 같다면 값을 result에 저장하고 종료한다.

만약 같지 않다면 root->right에 대해서 다시 재귀적으로 함수를 실행하면 된다.