在Binary搜索树中打印K个最小值

Question

I am trying to figure out how to print the lowest k values in a binary search tree. I an having trouble stopping the method

code:

def kthSmallestBST(node,k, count):
    if node == None or count == k:
        return
    else:
        kthSmallestBST(node.left, k, count)
        count += 1
        print node.data
        kthSmallestBST(node.right, k, count)
        count += 1


kthSmallestBST(BST, 3, 0)

Currently my output simply printed the whole tree inorder

Answer 1

Changes to the value of count don't propagate back up to the caller. You need to return the new count:

def kthSmallestBST(node,k, count):
    if node is None or count >= k:
        return 0
    else:
        count += kthSmallestBST(node.left, k, count)
        if count < k:
            print node.data
            count += 1
        count += kthSmallestBST(node.right, k, count)
        return count

Also note you don't need both k and count . You can get rid of count , decrement k instead of incrementing count , and compare k against 0 (instead of against count ). Here's what you get:

def kthSmallestBST(node, k):
    if node is None or k <= 0:
        return 0
    k = kthSmallestBST(node.left, k)
    if k > 0:
        print node.data
        k -= 1
    k = kthSmallestBST(node.right, k)
    return k

Answer 2

It's a rather "functional programming" solution, but one way is to generate (lazily) the nodes in the tree in order, and then using itertools to just take the first k.

def inorder(tree):
    if not tree: return
    for node in inorder(tree.left): yield node
    yield tree
    for node in inorder(tree.right): yield node

def least(tree, k):
    return itertools.islice(inorder(tree), k)

If you're using Python 3, you can use "yield from" to make this solution shorter.

Answer 3

You need to change things around a bit so you know how many elements were found during the recursive call. Have the function return the number of elements it found an add them up. Also you need to check the count between the recursive calls and the current node's element.

Something like:

def kthSmallestBST(node, k, count):
    if node == None or count == k:
        return 0
    else:
        count += kthSmallestBST(node.left, k, count)
        if(count == k) 
            return count
        print node.data
        count += 1
        count += kthSmallestBST(node.right, k, count)
        return count

Answer 4

import unittest

class BST(object):

    def __init__(self, key):
        self.key = key 
        self.left = None
        self.right = None


def get_kth_smallest_keys(node, k): 
    """Return, as a list, `k` smallest keys in a binary search tree rooted at `node`.

    """
    smallest = []
    _get_kth_smallest_keys(node, k, smallest)
    return smallest

def _get_kth_smallest_keys(node, k, smallest):
    """A helper function. Appends nodes to the given list, `smallest`, and stop
    when k reaches 0.

    Returns the number of nodes appended to said list.

    """
    if node is None or k == 0:
        return 0
    # first, recurse left, and we get the number of nodes appended to the list by that call
    nk = _get_kth_smallest_keys(node.left, k, smallest)
    # if that number already reduces our counter to zero, we fail fast, returning that same number
    if k - nk <= 0:
        return nk
    # otherwise, we can still append this node's key to the list
    smallest.append(node.key)
    # then we recurse right, with a counter that is less 1 (our append) and less nk (appended by the left recurse)
    nnk = _get_kth_smallest_keys(node.right, k - 1 - nk, smallest)
    # our return value is the sum of our append (1) and the appends from both recurse calls
    return nk + 1 + nnk 


class BSTTest(unittest.TestCase):

    def test_smallest_keys(self):
        root = BST(10)
        root.left = BST(6)
        root.right = BST(15)
        root.left.right = BST(8)
        root.right.right = BST(20)

        self.assertEquals(get_kth_smallest_keys(root, 0), []) 
        self.assertEquals(get_kth_smallest_keys(root, 1), [6])
        self.assertEquals(get_kth_smallest_keys(root, 3), [6, 8, 10])
        self.assertEquals(get_kth_smallest_keys(root, 5), [6, 8, 10, 15, 20])
        self.assertEquals(get_kth_smallest_keys(root, 6), [6, 8, 10, 15, 20])


if __name__ == '__main__':
    unittest.main()