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
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
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.
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
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()
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.