Search code examples
pythonrecursionbinary-search-treeexit

Python recursion - how to exit early


I've been playing with BST (binary search tree) and I'm wondering how to do an early exit. Following is the code I've written to find kth smallest. It recursively calls the child node's find_smallest_at_k, stack is just a list passed into the function to add all the elements in inorder. Currently this solution walks all the nodes inorder and then I have to select the kth item from "stack" outside this function.

def find_smallest_at_k(self, k, stack, i):
    if self is None:
        return i

    if (self.left is not None): 
        i = self.left.find_smallest_at_k(k, stack, i) 

    print(stack, i)
    stack.insert(i, self.data)
    i += 1
    if i == k: 
        print(stack[k - 1]) 
        print "Returning" 

    if (self.right is not None): 
        i = self.right.find_smallest_at_k(k, stack, i) 

    return i 

It's called like this,

our_stack = []
self.root.find_smallest_at_k(k, our_stack, 0) 
return our_stack[k-1] 

I'm not sure if it's possible to exit early from that function. If my k is say 1, I don't really have to walk all the nodes then find the first element. It also doesn't feel right to pass list from outside function - feels like passing pointers to a function in C. Could anyone suggest better alternatives than what I've done so far?


Solution

  • Passing list as arguments: Passing the list as argument can be good practice, if you make your function tail-recursive. Otherwise it's pointless. With BST where there are two potential recursive function calls to be done, it's a bit of a tall ask.

    Else you can just return the list. I don't see the necessity of variable i. Anyway if you absolutely need to return multiples values, you can always use tuples like this return i, stack and this i, stack = root.find_smallest_at_k(k).

    Fast-forwarding: For the fast-forwarding, note the right nodes of a BST parent node are always bigger than the parent. Thus if you descend the tree always on the right children, you'll end up with a growing sequence of values. Thus the first k values of that sequence are necessarily the smallest, so it's pointless to go right k times or more in a sequence.

    Even in the middle of you descend you go left at times, it's pointless to go more than k times on the right. The BST properties ensures that if you go right, ALL subsequent numbers below in the hierarchy will be greater than the parent. Thus going right k times or more is useless.

    Code: Here is a pseudo-python code quickly made. It's not tested.

    def findKSmallest( self, k, rightSteps=0 ):
        if rightSteps >= k: #We went right more than k times
            return []
        leftSmallest = self.left.findKSmallest( k, rightSteps ) if self.left != None else []
        rightSmallest = self.right.findKSmallest( k, rightSteps + 1 ) if self.right != None else []
        mySmallest = sorted( leftSmallest + [self.data] + rightSmallest )
        return mySmallest[:k]
    

    EDIT The other version, following my comment.

    def findKSmallest( self, k ):
        if k == 0:
            return []
        leftSmallest = self.left.findKSmallest( k ) if self.left != None else []
        rightSmallest = self.right.findKSmallest( k - 1 ) if self.right != None else []
        mySmallest = sorted( leftSmallest + [self.data] + rightSmallest )
        return mySmallest[:k]
    

    Note that if k==1, this is indeed the search of the smallest element. Any move to the right, will immediately returns [], which contributes to nothing.