Clean python k+log(n) code without stack


  • 0
    K

    We could use the same strategy as finding the closest value, using log(n) operations to find the ceil/floor. When we backtracking after find the ceil or floor, do a inorder / reverse inorder traversal until we get all k smaller / greater numbers or the node is null. After that, merge this two arrays in k operations. The whole time is ~ 3k + 2log(n) or O(k + log(n)).

    class Solution(object):
        def closestKValues(self, root, target, k):
            """
            :type root: TreeNode
            :type target: float
            :type k: int
            :rtype: List[int]
            """
            self.small_k = []
            self.great_k = []
    
            def visitL(node):
                if len(self.small_k) == k or not node: return
                visitL(node.right)
                if len(self.small_k) < k: self.small_k += [node.val]
                visitL(node.left)
    
            def floorK(node):
                if not node: return
                if node.val <= target:
                    floorK(node.right)
                    if len(self.small_k) < k: self.small_k += [node.val]
                    visitL(node.left)
                    return
                return floorK(node.left)
    
            def visitR(node):
                if len(self.great_k) == k or not node: return
                visitR(node.left)
                if len(self.great_k) < k: self.great_k += [node.val]
                visitR(node.right)
    
            def ceilK(node):
                if not node: return
                if node.val >= target:
                    ceilK(node.left)
                    if len(self.great_k) < k: self.great_k += [node.val]
                    visitR(node.right)
                    return
                return ceilK(node.right)
            
            floorK(root)
            ceilK(root)
    
            ans = []
            i, j = 0, 0
            for t in xrange(k):
                if i == len(self.small_k): 
                    ans += [self.great_k[j]]
                    j += 1
                elif j == len(self.great_k):
                    ans += [self.small_k[i]]
                    i += 1
                elif self.small_k[i] == self.great_k[j]:
                    ans += [self.small_k[i]]
                    i += 1
                    j += 1
                elif self.great_k[j] - target < target - self.small_k[i]:
                    ans += [self.great_k[j]]
                    j += 1
                else:
                    ans += [self.small_k[i]]
                    i += 1
            return ans

Log in to reply
 

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