[Java] Find k-th smallest element in BST

We saw how to find the second biggest element in a BST, now let's invert the search and look for the k-th smallest element instead.

We could do this operation in O(N) time by doing a preorder traversal of the tree and counting the elements visited so far, but this would not be taking advantage of the BST property. Since we know that, we can instead extend the Node definition to add a new value: leftSize. Each node now tracks the number of items in its left subtree, therefore allowing us to take advantage of the BST property while searching for the K-th smallest element.

The idea is as follows:

• the smallest element is K = 1 and will definitely sit all the way on the left
• nodes that do not have a left subtree have a leftSize of 0
• an element is exactly in position K only when its leftSize + 1 = K

Now, to use the BST property we need to decide where to redirect the search in case the property above does not hold for the current node. Imagining the tree nodes listed sequentially we can easily figure out that:

• if K <= leftSize of the current node, then we need to walk down the left subtree of this node. There are many nodes smaller than the current one, therefore we should search there.
• otherwise we have to walk down the right subtree of this node BUT we also need to adjust the value of K to K - 1 - leftSize of the current node before doing the recursive call on that subtree. This is because we know the current node (-1) and all the other nodes on the left subtree (-leftSize) can be discarded since they are in a position smaller than the requested K in the preorder traversal of the tree so we need to search in the remaining ones sitting on the right subtree. We basically convert the problem to finding the (K - already checked elements)-th smallest element in a new tree rooted at the first right child of the current node.

If none of the above conditions can be met, then K is bigger than the total tree size, therefore we should return null or throw an exception.

You can check my implementation of getKthSmallestElement on my Gist alongside some test cases in BSTJTests.

And bonus: Order statistic trees already implement this feature by design :)