Tony's Log

Algorithms, Distributed System, Machine Learning

  博客园 :: 首页 :: 博问 :: 闪存 :: 新随笔 :: 联系 :: 订阅 订阅 :: 管理 ::

Simply to utilize BST's property. The two-pass code below can be extended to support frequent insert\delete.

struct SNode
{
    SNode(int rcnt) : cnt(rcnt), pleft(nullptr), pright(nullptr){}
    int cnt;
    SNode *pleft;
    SNode *pright;
};
class Solution 
{
    SNode* build(TreeNode *p)
    {
        if (!p) return nullptr;

        SNode *pNew = new SNode(1);

        if (p->left)
        {
            pNew->pleft = build(p->left);
            pNew->cnt += pNew->pleft->cnt;
        }
        if (p->right)
        {
            pNew->pright = build(p->right);
            pNew->cnt += pNew->pright->cnt;
        }

        return pNew;
    }
public:
    int kthSmallest(TreeNode* root, int k) 
    {
        SNode *pS = build(root);
        int ret = 0;
        SNode *p = pS;
        TreeNode *pt = root;
        while (p)
        {
            int cnt = (p->pleft ? p->pleft->cnt : 0) + 1;
            if (cnt == k)
            {
                ret = pt->val;
                break;
            }
            if (cnt > k)
            {
                p = p->pleft;
                pt = pt->left;
            }
            if (cnt < k)
            {
                p = p->pright;
                pt = pt->right;
                k -= cnt;
            }
        }
        return ret;
    }
};

And this is a one-pass solution:

class Solution 
{
    int ret;
    int count(TreeNode *root, int k)
    {
        if (!root) return 0;
        int cntLeft = count(root->left, k);
        if (cntLeft == -1) return -1;
        if ((cntLeft + 1) == k)
        {
            ret = root->val;
            return -1;
        }

        int cntRight = count(root->right, k - cntLeft - 1);
        if (cntRight == -1) return -1;
return cntLeft + 1 + cntRight; } public: int kthSmallest(TreeNode* root, int k) { count(root, k); return ret; } };
posted on 2015-07-02 09:45  Tonix  阅读(143)  评论(0编辑  收藏  举报