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

Question

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

 

이진 탐색 트리와 정수 k가 주어졌을 때, 모든 노드 중 k번째로 작은 값을 리턴하는 함수를 작성하라.

 

제약사항

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

---

Solution

문제를 풀 때 고려한 사항

• root가 nullptr일 때 에도 정상 동작해야 한다. (제약사항을 보아하니 root가 비어있는 경우는 없다)
• 원소가 중복된 경우에는? 

---

Solution

무식하게 풀어본다. (Bruth Force)
모든 노드를 순환하여 k번째 작은 원소를 찾아준다. 
순환 순서를 inorder로 할 시 다시 정렬해줄 필요는 없다. 

class Solution {
public:
    void AddAllElement(vector<int> &node_val_vec, TreeNode* root){
        if(root == nullptr) return;
        
        if(root->left) AddAllElement(node_val_vec, root->left);
        node_val_vec.push_back(root->val);
        if(root->right) AddAllElement(node_val_vec, root->right);
    }
    int kthSmallest(TreeNode* root, int k) {
        vector<int> node_val_vec; 
        AddAllElement(node_val_vec, root);

        return node_val_vec[k-1];
    }
};

• Time complexity : O(N) 
• Space Complexity : O(N) 

해당 코드의 시간복잡도를 분석해보자면, AddAllElement 가 노드 개수만큼 실행되기 때문에 O(N)이다.
공간 복잡도도 vector를 위한 공간 O(N)이다.

Solution2

이진 탐색 트리의 특징을 잘 생각해본다. 
이진 탐색 트리는 루트 노드를 기준으로 왼쪽 자식 트리는 모두 루트 노드보다 값이 작고, 오른쪽 자식 트리는 모두 루트노드보다 값이 크다. 

즉, inorder 탐색을 진행한다면, k번째 탐색되는 노드가 k번째 작은 노드인 것을 알 수 있는데, 이 과정에서 모두를 탐색할 이유는 없다.

class Solution {
public:

    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode *> my_stack;

        while(true){
            while(root != nullptr){
                my_stack.push(root);
                root = root->left;
            }

            root = my_stack.top();
            my_stack.pop();
            if(--k == 0) return root->val;
            root = root->right;
        }
        
    }
};

• Time complexity : O(H + K) , H는 트리의 높이다
• Space Complexity : O(H) 

출처 : https://leetcode.com/explore/interview/card/top-interview-questions-medium/108/trees-and-graphs/790/

Explore - LeetCode