Augmented BST to resolve the question in the follow up

  • 0
    struct AugmentedTreeNode {
        int val;
        int index;
        AugmentedTreeNode *left;
        AugmentedTreeNode *right;
        AugmentedTreeNode(int x) : val(x),index(0),left(NULL),right(NULL) {}
    class Solution {
        AugmentedTreeNode* buildAugmentedTree(TreeNode* node) {
            if(!node) return nullptr;
            AugmentedTreeNode *root = new AugmentedTreeNode(node->val);
            root->left = buildAugmentedTree(node->left);
            root->right = buildAugmentedTree(node->right);
            return root;
        void setIndex(AugmentedTreeNode *node, int& prev) {
            if(!node) return;
            setIndex(node->left, prev);
            node->index = prev;
            setIndex(node->right, prev);
        int findk(AugmentedTreeNode *node, int k) {
            if(node->index==k) return node->val;
            else if(node->index < k) return findk(node->right, k);
            else if(node->index > k) return findk(node->left, k);
        int kthSmallest(TreeNode* root, int k) {
            int prev = 0;
            AugmentedTreeNode *augroot = buildAugmentedTree(root);
            setIndex(augroot, prev);
            return findk(augroot, k);

Log in to reply

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