O(h) complexity by add left subtree node count into TreeNode Structure


  • 3
    Y

    Modify TreeNode structure and add left subtree node count and find kth smallest element base on (http://www.geeksforgeeks.org/find-k-th-smallest-element-in-bst-order-statistics-in-bst/)

    The idea is to maintain rank of each node. We can keep track of elements in a subtree of any node while building the tree. Since we need K-th smallest element, we can maintain number of elements of left subtree in every node.

    Assume that the root is having N nodes in its left subtree. If K = N + 1, root is K-th node. If K < N, we will continue our search (recursion) for the Kth smallest element in the left subtree of root. If K > N + 1, we continue our search in the right subtree for the (K – N – 1)-th smallest element. Note that we need the count of elements in left subtree only.

    1.travel tree by level and insert node into TreeNodeWithCount Tree.

    2.find kth smallest in the TreeNodeWithCount Tree.

    public class TreeNodeWithCount {
        int val;
        int lCount;
        TreeNodeWithCount left;
        TreeNodeWithCount right;
        TreeNodeWithCount(int x) { val = x; }
    }
    
    public int kthSmallest(TreeNode root, int k) {
        if(root == null) return -1;
        TreeNodeWithCount rootWithCount = createBSTWithCount(root);
        return kthSmallestWithCount(rootWithCount, k);
    }
    
    public TreeNodeWithCount createBSTWithCount(TreeNode root) {
        Queue<TreeNode> queue = new LinkedList<TreeNode>();
        queue.add(root);
        TreeNodeWithCount rootWithCount = null;
        while(!queue.isEmpty()) {
            TreeNode node = queue.remove();
            TreeNodeWithCount nodeWithCount = new TreeNodeWithCount(node.val);
            rootWithCount = insertBSTWithCount(rootWithCount, nodeWithCount);
            if(node.left != null) queue.add(node.left);
            if(node.right != null) queue.add(node.right);
        }
        return rootWithCount;
    }
    
    public TreeNodeWithCount insertBSTWithCount(TreeNodeWithCount rootWithCount, TreeNodeWithCount nodeWithCount) {
        TreeNodeWithCount cur = rootWithCount, parent = rootWithCount;
        while(cur != null) {
            parent = cur;
            if(nodeWithCount.val < cur.val) {
                cur.lCount++;
                cur = cur.left;
            } else {
                cur = cur.right;
            }
        }
        if(rootWithCount == null) {
            rootWithCount = nodeWithCount;
        } else if(nodeWithCount.val < parent.val) {
            parent.left = nodeWithCount;
        } else {
            parent.right = nodeWithCount;
        }
        return rootWithCount;
    }
    
    public int kthSmallestWithCount(TreeNodeWithCount rootWithCount, int k) {
        while(rootWithCount != null) {
            if(k == rootWithCount.lCount + 1) {
                return rootWithCount.val;
            } else if(k <= rootWithCount.lCount) {
                rootWithCount = rootWithCount.left;
            } else {
                k = k - rootWithCount.lCount - 1;
                rootWithCount = rootWithCount.right;
            }
        }
        return -1;
    }

  • 0
    Z

    When you traverse this tree in order to create TreeNodeWithCount, I guess you already can figure out what the kth smallest number is, so why bother to give step.


  • 0
    A

    When the BST is modified often and we need to find Kth element frequently, the pre-process will be worthy.


  • 0
    O

    "we can maintain number of elements of left subtree in every node"

    Than's not intuitive at all. Why not maintain a number of elements of subtree rooted with every node. Take a look at this discussion:

    https://leetcode.com/discuss/87612/why-the-hint-says-the-optimal-runtime-complexity-height-bst


Log in to reply
 

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