Short 75ms python


  • 0
    M

    The thinking is simple:

    1. Find top k numbers that are less or equal than target.
    2. Find top k numbers that are greater than target.
    3. Merge, return first k elements.

    I wrote two functions at first:

    def lte(res, root, target, k):
        if not root or len(res) >= k:
            return
        if root.val <= target:
            lte(res, root.right, target, k)
            if len(res) >= k:
                return
            res.append(root.val)
        lte(res, root.left, target, k)
    
    def gt(res, root, target, k):
        if not root or len(res) >= k:
            return
        if root.val > target:
            gt(res, root.left, target, k)
            if len(res) >= k:
                return
            res.append(root.val)
        gt(res, root.right, target, k)
    

    This is clear and easy to understand.
    However I combine them together and used lambda, just for short and beauty.

    class Solution(object):
        def closestKValues(self, root, target, k):
                
            def inorder(res, root, target, k, lte):
                if not root or len(res) >= k:
                    return
                if (root.val <= target) == lte:
                    inorder(res, root.right if lte else root.left, target, k, lte)
                    if len(res) >= k:
                        return
                    res.append(root.val)
                inorder(res, root.left if lte else root.right, target, k, lte)
    
            lte, gt = [], []
            inorder(lte, root, target, k, True)
            inorder(gt, root, target, k, False)
            res = lte + gt
            res.sort(cmp=lambda x, y:cmp(abs(x - target), abs(y - target)) )
            return res[:k]
            
    
    

Log in to reply
 

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