Two 8 lines solutions in C++, easy to understand


  • 0

    Solution 1

    Count nodes of left subtree, if it equals k, we found the kth smallest element.

    class Solution {
    public:
        int DFS(TreeNode* root, int k, int& val){
           if(!root) return 0;
           int left = DFS(root->left, k, val) + 1;
           int right = DFS(root->right, k - left, val) + 1;
           if(k == left) val = root->val;
           return left + right - 1;
        }
        
        int kthSmallest(TreeNode* root, int k) {
            int val = 0;
            DFS(root, k, val);
            return val;
        }
    };
    

    If you want to stop as soon as we found the kth smallest element, you can add a bool to DFS function, it reduces run time from 16 ms to 12 ms and beats 97.13% of C++ solution.

    Code:

    class Solution {
    public:
        int DFS(TreeNode* root, int k, int& val, bool& found){
           if(!root || found) return 0;
           int left = DFS(root->left, k, val, found) + 1;
           int right = DFS(root->right, k - left, val, found) + 1;
           if(k == left) val = root->val, found = true;
           return left + right - 1;
        }
        
        int kthSmallest(TreeNode* root, int k) {
            int val = 0;
            bool found = false;
            DFS(root, k, val, found);
            return val;
        }
    };
    

    Solution 2

    Using stack, push right subtree first, then root, then right subtree, it's a reversed in-order traversal.

    class Solution {
    public:
        int kthSmallest(TreeNode* root, int k) {
            stack<int>s;
            reversedInOrder(root, s);
            while(--k) s.pop();
            return s.top();
        }
        
        void reversedInOrder(TreeNode* root, stack<int>& s){
            if(!root) return;
            reversedInOrder(root->right, s);
            s.push(root->val);
            reversedInOrder(root->left, s);
        }
    };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.