Closest Binary Search Tree Value II

Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.

Note:

• Given target value is a floating point.
• You may assume k is always valid, that is: k ≤ total nodes.
• You are guaranteed to have only one unique set of k values in the BST that are closest to the target.

Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

这是Closest Binary Search Tree Value的一道拓展题目。查找和target最接近的k个值。

思路其实是对BST中最接近target的值的前继和后继做一个merge。合并出k个值即可。我的直接思路是先找到最近的这个节点，之后每次调用getPredecessor(N)和getSuccessor(N)的函数来获得当前处理到的结点的前继或者后继。找到最靠近的结点的复杂度为O(logn)(Closest Binary Search Tree Value),之后每次查找前继和后继的复杂度为O(logn)(Inorder Successor in BST)。所以总的复杂度为O(logn+ klogn) = O(klogn)。代码如下：

class Solution(object):
    def closestKValues(self, root, target, k):
        """
        :type root: TreeNode
        :type target: float
        :type k: int
        :rtype: List[int]
        """
        #first find the closet value, then find the predecessor or successor of the node, merge k
        closet = root
        cur = root
        while cur:
            if abs(cur.val - target) < abs(target - closet.val):
                closet = cur
            if cur.val > target:
                cur = cur.left
            else:
                cur = cur.right
        pre = self.getPredecessor(root, closet)
        succ = self.getSuccessor(root, closet)
        res = [closet.val]
        k -= 1
        while k and pre and succ:
            if abs(pre.val - target) < abs(succ.val - target):
                res.append(pre.val)
                pre = self.getPredecessor(root, pre)
            else:
                res.append(succ.val)
                succ = self.getSuccessor(root, succ)
            k -= 1
     
        if k > 0:
            if pre:
                while k:
                    res.append(pre.val)
                    pre = self.getPredecessor(root, pre)
                    k -= 1
            else:
                while k:
                    res.append(succ.val)
                    succ = self.getSuccessor(root, succ)
                    k -= 1
        return res
        
    def getPredecessor(self, root, p):
        pre = None
        while root:
            if root.val < p.val:
                pre = root
                root = root.right
            else:
                root = root.left
        return pre
        
    def getSuccessor(self, root, p):
        succ = None
        while root:
            if root.val > p.val:
                succ = root
                root = root.left
            else:
                root = root.right
        return succ

但是这种解法显然不是最优的，每次查找前继和后继都是从树根开始的，中间每次重复走了很多路程。一个比较好的思路就是类似于Binary Search Tree Iterator，利用一个栈存储元素，每次pop元素时可以压入后继的序列。保存整个查找路径的值，节约了很多的时间。代码如下：

# 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]
        """
        res = []
        pred = []
        succ = []
        cur = root
        while cur:
            if cur.val > target:
                succ.append(cur)
                cur = cur.left
            else:
                pred.append(cur)
                cur = cur.right
                
        while k:
            if not pred and not succ:
                break
            elif not pred:
                res.append(self.getSuccessor(succ))
            elif not succ:
                res.append(self.getPredecessor(pred))
            elif abs(pred[-1].val - target) < abs(succ[-1].val - target):
                res.append(self.getPredecessor(pred))
            else:
                res.append(self.getSuccessor(succ))
            k -= 1
                
        return res
            
    def getPredecessor(self, st):
        popped = st.pop()
        p = popped.left
        while p:
            st.append(p)
            p = p.right
        return popped.val
        
    def getSuccessor(self, st):
        popped = st.pop()
        p = popped.right
        while p:
            st.append(p)
            p = p.left
        return popped.val
        
        

这种解法的时间和空间复杂度理论上都是O(klogn)。但是时间复杂度显然要优于我想出的第一个解法。