Recursively finding the kth largest node in a BST

Question

I'm doing some self-studying right now and came across a really mysterious piece of code that "finds the Kth largest node in a BST"

       4           
      / \  
     2   8 
        / \  
       5  10                  

K = 2, Output = 8

Here's the code:

public TreeNode findKthLargest(TreeNode root, int k) {                   
        if(root == null) return null;
        int rightSize=0;
        if(root.right != null) {
            rightSize=size(root.right);
        }
        if(rightSize+1 == k) {
            return root;      
        } else if(k <= rightSize) {
            return findKthLargest(root.right, k);
        } else { 
            return findKthLargest(root.left, k-rightSize-1);
        }
    }

Can anyone give me a hint or point me in the right direction on HOW this code works? It looks utterly mysterious to me.

Why are we paying so much attention to the "right" side of the BST? Why do we completely ignore the left side?

What is the point of using "rightSize+1 == k" as our check to return our answer?

Note:

Upon second inspection, the "size" function is really weird. I'm practicing on Firecode.io right now and this is their provided structures:

class TreeNode {
    int data;
    TreeNode left;
    TreeNode right;
    
    TreeNode() {
    }

    TreeNode(int data) {
        this.data = data;
    }

    TreeNode(int data, TreeNode left, TreeNode right) {
        this.data = data;
        this.left = left;
        this.right = right;
    }
}   

public int size(TreeNode root);

The 'size' function at the end is not elaborated upon so I'm unsure if this is just some...error on Firecode's part...

Upon second inspection, the "size" function is really weird -- do you have the definition of this function? I assume it's a count of all nodes in the parameter subtree? — ggorlen
I am also assuming that, but FireCode doesn't provide the full function. So not sure what it does exactly. From spending a day mulling over the code, it looks like we're doing some sort of reverse inorder traversal where I add the right-most node, decrement k, add right-most node, decrement k...until it hits 1 or 0. But this is just a guess — buddybuddybuddybuddy

trincot trincot · Accepted Answer · 2021-03-09T16:08:02

We may assume that the size method returns an integer that represents the number of nodes in the subtree rooted by the node given as argument. So in your example tree:

...if we would pass the root node (with value 4), size will return 5. If we pass the right-child of the root (with value 8), size will return 3. For any other node in this example it will return 1.

Let's see what happens with different values of K. I will for simplicity assume that the tree has no duplicate values (but the reasoning would remain the same).

K=4

Say that K=4, then the answer would have to be 4 (the value of the root node). So first the code performs this:

rightSize=size(root.right);

This will return 3. Then we have the following comparison:

if(rightSize+1 == k) {

The reasoning here is that all values that are greater than the root's value, and only those values, are necessarily located at the right side of the root. So if we know the count of nodes in the right subtree, then we know how many values are greater than the root's value. By consequence, the root has one more as "reverse-order" number.

In the example, the root has 3 nodes at the right, so the root has the 4^th-greatest number.

K=2

With K=2 we still find the size of the root's right subtree is 3, and so our current value is too small to be the 2nd largest value. Therefore we must descend to the right. This is the reason to make the following recursive call:

return findKthLargest(root.right, k);

In that recursive call we arrive at a similar situation as in the first described case: the root (which is now the node with value 8), has a right subtree of size 1. And so the root itself has the 2^nd greatest value, which is the solution.

K=3

Again we find the root's right subtree has a size of 3, and we also make that recursive call:

return findKthLargest(root.right, k);

In that recursive call we learn again that the right subtree (of 8) has size 1. And now we get into the third case:

return findKthLargest(root.left, k-rightSize-1);

Here we move left, as we know the node to find is not at the right, and is not the current node (8)... so it must be at the left. We should not forget that we skip here the right subtree, which has values that are greater. And also we skip the current node (8), which is also greater than any value in the left subtree. So to have a correct situation, we must subtract that number of nodes (both rightSize and 1) from the value of K that we will pass on to the recursive call, where we will have no clue anymore about the existence of those nodes (every recursive call treats the given node as a root).

In that deeper recursive call (at node 5) we thus have K=1. Its right subtree has size 0 (as there is none), and so we can conclude that this node (with 5) is indeed the greatest value in this "tree" (it happens to be a single node). So we return it as the solution.

Backtracking

Note that the value that a recursive call returns, is immediately returned up the recursion tree as-is. So once a particular node is returned (in the first if block), you know that it is going to be that node that the original call will receive.

Recursively finding the kth largest node in a BST

1 Answers

K=4

K=2

K=3

Backtracking