Uber | Technical Phone Screen | kth largest element in BST - follow up - space O(1)
Anonymous User
10738

1 hr coding round.
I was asked a variation of this problem. Problem was to find kth largest instead of kth smallest element.
https://leetcode.com/problems/kth-smallest-element-in-a-bst/
'''
Given a Binary Search Tree (BST) and a positive integer k, find the k’th largest element in the Binary Search Tree.
For example, in the following BST, if k = 3, then output should be 15, and if k = 5, then output should be 4.
10
/ \
4 20
/ / \
2 15 40

'''
I quickly wrote a O(N) time and O(height) space solution as is in the solution for the problem above.

The interviewer then asked the follow up question - can we do it in O(1) space.
We discussed a few ideas and eventually I picked up on his hints to populate the left pointers of the leaf nodes to appropriate parents. Hence, we won't need a stack to store the elements and eliminate the O(N) space by storing the pointers. When there is no right element, it means we are at kth largest element currently, and subsequently we decrease k each time there's no right element.

So - populate Node(40)'s left to Node(20), and Node(15)'s left to Node(10). To do this, we'll need to utilize the property of BST. The next element in sorted array would be the leftmost element in the right subtree of that tree.

Spent some time after the interview to refine my answer and tested it. Here's the code with O(N) time and O(1) space. Note: Interviewer mentioned that modifying the tree structure is okay, but adding extra properties to Tree Node is not (so can't add parent property or something similar to the Node).

def find_kth_largest(root, k):
    if not root: return -1
    while k>0:
        start, end = root, root.right
        # If no right, it means we are at the largest element currently
        if not end:
            # Decrease k by 1
            k-=1
            # Point root to root.left
            root = root.left
            continue
            
        # Continually go to left of the right sub-tree
        # to find the next element which is greater than root
        # Property of BST
        while end.left:
            end = end.left
            
        # Re-wire connection to next element in reverse order
        end.left = start
        
        # Change root.right to None since right side is already traversed now
        root, root.right = root.right, None
    return start
n = find_kth_largest(root, 2)

EDIT: Made it to onsite!

Comments (11)