Explaining how to use Morris Traversal to find the k'th smallest element in a binary search tree

morris-traversal-kth-smallest-in-bst.py

# Problem:
# Given a binary search tree t with each node having properties left, right 
# (the left and right subtrees respectively) and value (the value of that 
# node) and an integer k, find the k-th smallest element in the BST using
# constant space.
#
# Example:
# t =
#    3
#  /   \
# 1     5
#      / \
#     4   6
#
# kthSmallestInBST(t, 4) should give 5; the elements in order are
# [1, 3, 4, 5, 6], and k is 1-indexed, so the 4th element is 5.
#
# Approach:
# What we want to do is an in order tree traversal: for a node t, first visit
# all of t.left's subtree, then t.value, then t.right's subtree.
#        
# The problem is that since we need to visit t.left's subtree first, we can't
# get back to t unless we store a pointer back to it. But this problem exists
# at each level of the tree: storing a pointer back to t will need to be done
# for each left subtree of t's right and left subtrees.
#
# So we can't do that, because it will use space proportional to the height
# of the tree, and likewise using recursion to stash that will use stack 
# space proportional to the height of the tree.
#
# Instead, what we can do is use a pointer in the tree itself to point back 
# to t. We have fundamentally 2 ways we could store this:
#
# * replace t.left.value value with a tuple of (t.left.value, t)
# * find a node with an unused left or right link, and use it to store the
#    back pointer
#
# Of these 2 solutions, the neater one is the second option, because it turns
# out that in an in-order traversal, after you visit the right-most child of
# the left subtree of t, then the next step is to visit t itself - and by
# definition, the right-most child is a leaf node, so it has a spare right
# link, which can be used to store a pointer back to t.
# 
# This is known as a Morris Traversal.
def kthSmallestInBST(t, k):
    while t is not None:
        # The problem of storing a pointer back to t only exists if t has a
        # left subtree, so let's first handle the case where we don't have a
        # left subtree to worry about.
        if t.left is None:
            if k == 1:
                return t.value
            k -= 1
            t = t.right
        else:
            # Now since we know we have a left subtree, before we look at 
            # that left subtree, we need to find the rightmost child of the
            # left subtree of t.
            # While doing this, we need to bear in mind that the right-most
            # child of t's left subtree might already point back at t - if
            # find this, then it must mean that it's time to consider t's
            # value itself.
            pre = t.left
            while pre.right is not None and pre.right is not t:
                pre = pre.right
            
            if pre.right is None:
                # pre is now the rightmost child of t's left subtree, and we know
                # that its right child link is empty - so we can temporarily 
                # hijack that link to point back at t, which prevents us from 
                # having to store a separate list of pointers back to parent nodes.
                pre.right = t
                t = t.left
            else:
                # We found that the rightmost child of the left subtree is pointing
                # at t already! So revert the change we made before, and consider
                # t.value, because this means we have processed all the nodes in 
                # the left subtree of t, and we can now move to t.right's subtree.
                pre.right = None
                if k == 1:
                    return t.value
                k -= 1
                t = t.right

This is a great write-up on how to use Morris to solve k smallest. I think it's more clear for me about how you can traverse a tree without recursion or stacks while using O(1) space.