# Two solutions, O(n) and O(log(n)+k)

• Both took 76ms. Please forgive me for the rather unpolished codes, my baby is crying...

The O(n) solution is really intuitive, easy to envision. Since this is BST, we can make it a sorted list, and finding the k nearest in a sorted list is easy. A little bit of trick is that, you do not have to go through the whole tree. Do an inorder traverse, maintain a k-sized moving window once your list grows larger than k, and once the new element is farther from target than the first element in your window, you can stop growing the list, and return the window as your result. I have the code here:

``````class Solution(object):
def closestKValues(self, root, target, k):
"""
:type root: TreeNode
:type target: float
:type k: int
:rtype: List[int]
"""
stack=[]
values=[]

probe=root
self.p1=self.p2=0

def kregion(v):
values.append(v)
if len(values)==k:
self.p2=k-1
return False
elif len(values)>k:
if abs(values[-1]-target)<abs(values[self.p1]-target):
self.p1+=1
self.p2+=1
return False
else:
return True

while probe:
if probe.left:
stack.append(probe)
probe=probe.left
else:
if kregion(probe.val):
break
while not probe.right:
if stack:
probe=stack.pop()
if kregion(probe.val):
break
else:
break
if not probe.right:
break
else:
probe=probe.right

return values[self.p1:self.p2+1]
``````

The O(log(n)+k) way is also pretty straight forward. First you find the insertion point for the target value in the tree (of course you do not actually insert it). This way, you get the nearest smaller and larger elements to the target (let's call them left and right). Even better, during the path to that point, you can put larger elements in a rightStack, and smaller elements in leftStack. Using these two stacks, you can do in-order traversal to find the immediately larger elements, and reverse in-order traversal for the immediately smaller elements, and of course you stop once you get a total of k values. Since we know that in-order and reverse both could take O(k) time, and initial location finding is O(log(n)), the total time is O(log(n)+k).

``````class Solution(object):
def closestKValues(self, root, target, k):
"""
:type root: TreeNode
:type target: float
:type k: int
:rtype: List[int]
"""
self.rightStack=[]
self.leftStack=[]
self.left=self.right=targetNode=None

result=[]
probe=root

while probe:
if probe.val > target:
self.rightStack.append(probe)
probe=probe.left
elif probe.val < target:
self.leftStack.append(probe)
probe=probe.right
else:
targetNode=probe
break

if targetNode:
result.append(targetNode.val)
self.left=self.right=targetNode
self.nextLeft()
self.nextRight()

else:
self.left=self.leftStack.pop() if self.leftStack else None
self.right=self.rightStack.pop() if self.rightStack else None

lv=self.nextLeft() if self.left else -float('inf')
rv=self.nextRight() if self.right else float('inf')

while len(result)<k:
if abs(lv-target)<abs(rv-target):
result.append(lv)
lv=self.nextLeft() if self.left else -float('inf')

else:
result.append(rv)
rv=self.nextRight() if self.right else float('inf')

return result

def nextLeft(self):
value=self.left.val

if self.left.left:
self.left=self.left.left
while self.left.right:
self.leftStack.append(self.left)
self.left=self.left.right
elif self.leftStack:
self.left=self.leftStack.pop()
else:
self.left=None
return value

def nextRight(self):
value=self.right.val

if self.right.right:
self.right=self.right.right
while self.right.left:
self.rightStack.append(self.right)
self.right=self.right.left
elif self.rightStack:
self.right=self.rightStack.pop()
else:
self.right=None
return value``````

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