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


  • 0
    C

    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

Log in to reply
 

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