Problem statement

https://leetcode.com/problems/closest-binary-search-tree-value-ii/

Solution

The idea is to keep two stacks: one is stack of predecessors and one is stack of successors. Then we $k$ times chose the best element of these two stack — closest to target, remove it from one of the stacks and instead of it, add to stack new candidates, amortized time for this is $O(1)$.

I need to explain more about how stacks of predecessors and successors works.

Complexity

Time complexity is $O(k + \log n)$ for balanced tree, $k$ is to work with stacks and $\log n$ to fill stacks initially.

Code

class Solution:
    def closestKValues(self, root, target, k):
        def goUpperNext(node):
            while node:
                upperStack.append(node)
                node = node.left

        def goLowerNext(node):
            while node:
                lowerStack.append(node)
                node = node.right
         
        ans, lowerStack, upperStack, p = [], [], [], root
 
        while p:
            if p.val < target:
                lowerStack.append(p)
                p = p.right
            else:
                upperStack.append(p)
                p = p.left

        for i in range(k):
            diff1, diff2 = None, None
           
            if upperStack:
                diff1 = upperStack[-1].val - target
            if lowerStack:
                diff2 = target - lowerStack[-1].val

            if not lowerStack or (diff1 != None and diff2 != None and diff1 <= diff2):
                top = upperStack.pop()
                goUpperNext(top.right)
            else:
                top = lowerStack.pop()
                goLowerNext(top.left)
                    
            ans.append(top.val)
        return ans