# Python Solution using heapq and bfs (simple to understand)

• Below is my code, including my tests. I just visit each node of the tree and add it to a heap. The heap is sorted according to distance to target. I avoid searching the tree for sizes larger than those popped off the heap. Most heap operations are O(log(n)) and thus the bottleneck is traversing the tree. Since the traversal is being pruned by max_size this should reduce the search space below O(n).

``````# Definition for a binary tree node.
class TreeNode(object):
def __init__(self, x):
self.val = x
self.left = None
self.right = None

def build_tree(arr):
if len(arr) == 0:
return None
root = TreeNode(arr[0])
q = [root]
while arr:
t = q.pop(0)
if arr:
if arr[0] == None:
arr.pop(0)
else:
t.left = TreeNode(arr.pop(0))
q.append(t.left)
if arr:
if arr[0] == None:
arr.pop(0)
else:
t.right = TreeNode(arr.pop(0))
q.append(t.right)
return root

class hitem:
def __init__(self, size, x):
self.val = x
self.size = size
def __lt__(self, other):
return self.size > other.size

class Solution(object):
def closestKValues(self, root, target, k):
"""
:type root: TreeNode
:type target: float
:type k: int
:rtype: List[int]
"""
from heapq import heappop, heappush
from sys import maxsize
h = []
q = [root]
max_size = maxsize
while q:
n = q.pop(0)
size = abs(target - n.val)
if size > max_size:
continue
heappush(h, hitem(size, n.val))
while len(h) > k:
item = heappop(h)
max_size = min(max_size, item.size)

for child in [n.left, n.right]:
if child:
q.append(child)

while len(h) > k:
heappop(h)

return [x.val for x in h]

def main():
s = Solution()
tests = [
(([3,1,4,None,2], 2.000000, 1), [2]),
(([3,2,4,1], 2.000000, 3), [2,1,3]),
]
for t in tests:
input_, output = t
print()
bst_arr, target, k = input_
ret = s.closestKValues(build_tree(bst_arr), target, k)
print(ret)
if set(ret) == set(output):
print("PASS")
else:
print("FAIL")
print("expected:")
print(output)
print("got:")
print(ret)

if __name__ == '__main__':
main()
``````

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.