牛客OJ：二叉搜索树第K个节点

/*
struct TreeNode {
    int val;
    struct TreeNode *left;
    struct TreeNode *right;
    TreeNode(int x) :
            val(x), left(NULL), right(NULL) {
    }
};
*/
struct MyTreeNode{
    TreeNode* p;
    MyTreeNode* left;
    MyTreeNode* right;
    int sumL;
    int sumR;
};

class Solution {
public:
    MyTreeNode* presolve(TreeNode* p,int& k){
        if(p == NULL) return NULL;
        MyTreeNode* res = new MyTreeNode();
        res->p = p;
        int sum1 = 0;
        int sum2 = 0;
        res->left = presolve(p->left,sum1);
        res->right = presolve(p->right,sum2);
        res->sumL = sum1;
        res->sumR = sum2;
        k = sum1 + sum2 + 1;
        return res;
    }
    TreeNode* findit(MyTreeNode* p,int k){
        if(p == NULL) return NULL;
        if(k > p->sumL + p->sumR + 1) return NULL;
        if(k == p->sumL + 1) return p->p;
        if(k <= p->sumL) return findit(p->left,k);
        if(k > p->sumL + 1) return findit(p->right,k-p->sumL-1);
        return NULL;
    }
    TreeNode* KthNode(TreeNode* pRoot, int k)
    {
        int sum = 0;
        MyTreeNode* p = presolve(pRoot,sum);
        return findit(p,k);
    }

    
};