Kth Smallest Element in a BST - Java Code

To calculate the k^th smallest element, we use the following property of a BST that is:

• The left sub-tree of a node T contains the nodes whose keys ( values ) are less than that of the current node T.
• The right sub-tree of a node T contains the nodes whose keys ( values ) are greater than that of the current node T.

So, based on the above two properties, we can think of the following algorithm to find out kth smallest element.

Algorithm

For a node T,

• If k = num_nodes(left Subtree of T) + 1, then the answer is the value in node T
• If k < num_nodes(left Subtree of T) + 1, then the k^th smallest is in the left sub-tree, so we reduce the problem to finding the k^th smallest element in the left subtree.
• If k > num_nodes(left Subtree of T) + 1, then the k^th smallest is in the right sub-tree, So, we reduce the problem to finding the { k - num_nodes(left subtree of T) - 1 } smallest element of the right subtree.

Java Code

class TreeNode {

       int val;

       TreeNode left;

       TreeNode right;

       TreeNode(int x) {

              val = x;

       }

}

public int kthSmallest(TreeNode root, int k) {

       int leftSubtreeCount = getNoOfNodes(root.left) + 1; // one for the current node

       if (k == leftSubtreeCount) {

              return root.val;

       } else if (k < leftSubtreeCount) {

              return kthSmallest(root.left, k);

       } else {

              return kthSmallest(root.right, k - leftSubtreeCount);

       }

}

private int getNoOfNodes(TreeNode currentNode) {

       if (currentNode == null) {

              return 0;

       }

       return 1 + getNoOfNodes(currentNode.left)

                     + getNoOfNodes(currentNode.right);

}