# Time:  O(h + k)
# Space: O(h)

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
    
class Solution(object):
    def closestKValues(self, root, target, k):
        """
        :type root: TreeNode
        :type target: float
        :type k: int
        :rtype: List[int]
        """
        stack = []
        node = self.closestValue(root, target, stack)
        result = [node.val]
        
        smaller_it, larger_it = BSTIterator(node, stack), BSTIterator(node, stack)
        smaller, larger = smaller_it.prev(), larger_it.next()
        while len(result) < k:
            if abs(smaller - target) < abs(larger - target):
                result.append(smaller)
                smaller = smaller_it.prev()
            else:
                result.append(larger)
                larger = larger_it.next()
        return result

    def closestValue(self, root, target, stack):
        """
        :type root: TreeNode
        :type target: float
        :rtype: int
        """
        gap = float("inf")
        closet = float("inf")
        while root:
            stack.append(root)
            if abs(root.val - target) < gap:
                gap = abs(root.val - target)
                closet = root
            if target == root.val:
                closet = stack.pop()
                break
            elif target < root.val:
                root = root.left
            else:
                root = root.right
        return closet


class BSTIterator:
    # @param root, a binary search tree's root node
    def __init__(self, cur, stack):
        self.stack = list(stack)
        self.cur = cur

    # @return an integer, the next number
    def next(self):
        node = None
        if self.cur and self.cur.right:
            node = self.cur.right
            while node.left:
                self.stack.append(node)
                node = node.left
        elif self.stack:
            node = self.stack.pop();
            while node:
                if node.val > self.cur.val:
                   break
                else:
                   node = self.stack.pop() if self.stack else None
        self.cur = node
        return node.val if node else float("inf")
    
    # @return an integer, the previous number
    def prev(self):
        node = None
        if self.cur and self.cur.left:
            node = self.cur.left
            while node.right:
                self.stack.append(node)
                node = node.right
        elif self.stack:
            node = self.stack.pop();
            while node:
                if node.val < self.cur.val:
                   break
                else:
                   node = self.stack.pop() if self.stack else None
        self.cur = node
        return node.val if node else float("-inf")

@@ -422,31 +422,24 @@
 = %5B%5D%0A       
- node =
  self.closes
@@ -484,20 +484,25 @@
 sult = %5B
-node
+stack%5B-1%5D
 .val%5D%0A  
@@ -548,30 +548,24 @@
 BSTIterator(
-node, 
 stack), BSTI
@@ -576,14 +576,8 @@
 tor(
-node, 
 stac
@@ -1370,50 +1370,14 @@
     
-closet = stack.pop()%0A                break
+return
 %0A   
@@ -1506,22 +1506,68 @@
     
-return closet%0A
+while stack and stack%5B-1%5D != closet:%0A            stack.pop()
 %0A%0Acl
@@ -1661,13 +1661,8 @@
 elf,
- cur,
  sta
@@ -1718,19 +1718,32 @@
 f.cur = 
-cur
+self.stack.pop()
 %0A%0A    # 
@@ -1850,32 +1850,72 @@
 self.cur.right:%0A
+            self.stack.append(self.cur)%0A
             node
@@ -2051,32 +2051,60 @@
 lif self.stack:%0A
+            prev = self.cur%0A
             node
@@ -2176,28 +2176,23 @@
 ode.
-val %3E self.cur.val
+left is prev
 :%0A
+ 
     
@@ -2226,32 +2226,65 @@
           else:%0A
+                    prev = node%0A 
                 
@@ -2534,24 +2534,64 @@
 f.cur.left:%0A
+            self.stack.append(self.cur)%0A
             
@@ -2732,32 +2732,60 @@
 lif self.stack:%0A
+            prev = self.cur%0A
             node
@@ -2857,28 +2857,24 @@
 ode.
-val %3C self.cur.val
+right is prev
 :%0A
+ 
     
@@ -2908,32 +2908,65 @@
           else:%0A
+                    prev = node%0A 
                 

Update closest-binary-search-tree-value-ii.py