A/C Python solution, use heapq, easy to understand, beat 5%


  • 0
    W
    def kSmallestPairs(self, nums1, nums2, k):
    
    
        # print "nums1 = ", nums1
        # print "nums2 = ", nums2
        # print "k = ", k
    
        import heapq
        class local:
            sum_idx = []
            res = []
            resVal = []
    
        if nums1 == [] or nums2 == []:
            return []
        heapq.heappush(local.sum_idx, [nums1[0]+nums2[0], [0,0]])
        # heapq.heappush(local.sum_idx, [nums1[0]+nums2[0], [-1,0]])
    
        # heapq.heappush(local.sum_idx, [2, [[1,3]]])
        # heapq.heappush(local.sum_idx, [1, [[0,0]]])
    
        # tmpVal = heapq.heappop(local.sum_idx)
        # print "tmpVal = ", tmpVal
    
        def validIndex(idx1, idx2):
            #print "idx1 = ", idx1, " idx2 = ", idx2
            if idx1 >=0 and idx1 < len(nums1) and idx2 >=0 and idx2 < len(nums2):
                #print "true"
                return True
            else:
                #print "false"
                return False
    
        def helper(k, curNum):
            #print "helper called, len(local.sum_idx) = ", len(local.sum_idx)
            #print "helper k = ", k, "local.resVal = ", local.resVal
            if len(local.sum_idx) == 0:
                return
    
            nextVal = heapq.heappop(local.sum_idx)
            #print "nextVal = ", nextVal
            curIdx1 = nextVal[1][0]
            curIdx2 = nextVal[1][1]
            #print "curIdx1 = ", curIdx1, " curIdx2 = ", curIdx2
    
            if curNum ==k:
                return
    
            local.res.append([curIdx1, curIdx2])
            local.resVal.append([nums1[curIdx1], nums2[curIdx2]])
    
    
    
            if validIndex(curIdx1-1, curIdx2+1) and [curIdx1-1, curIdx2+1] not in local.res: #3
                nextIdx1 = curIdx1 - 1
                nextIdx2 = curIdx2 + 1
    
                heapq.heappush(local.sum_idx, [nums1[nextIdx1] + nums2[nextIdx2], [nextIdx1, nextIdx2]])
                local.res.append([nextIdx1, nextIdx2])
            if validIndex(curIdx1, curIdx2+1) and [curIdx1, curIdx2+1] not in local.res: #6
                nextIdx1 = curIdx1
                nextIdx2 = curIdx2 +1
    
                heapq.heappush(local.sum_idx, [nums1[nextIdx1] + nums2[nextIdx2], [nextIdx1, nextIdx2]])
                local.res.append([nextIdx1, nextIdx2])
    
            if validIndex(curIdx1 + 1, curIdx2 -1) and [curIdx1 + 1, curIdx2-1] not in local.res: #7
                nextIdx1 = curIdx1 + 1
                nextIdx2 = curIdx2 - 1
    
                heapq.heappush(local.sum_idx, [nums1[nextIdx1] + nums2[nextIdx2], [nextIdx1, nextIdx2]])
                local.res.append([nextIdx1, nextIdx2])
    
            if validIndex(curIdx1 + 1, curIdx2) and [curIdx1 + 1, curIdx2] not in local.res: #8
                nextIdx1 = curIdx1 + 1
                nextIdx2 = curIdx2
    
                #print "nextIdx1 = ", nextIdx1, " nextIdx2 = ", nextIdx1
                heapq.heappush(local.sum_idx, [nums1[nextIdx1] + nums2[nextIdx2], [nextIdx1, nextIdx2]])
                local.res.append([nextIdx1, nextIdx2])
    
            if validIndex(curIdx1 + 1, curIdx2 + 1) and [curIdx1 + 1, curIdx2 + 1] not in local.res: #9
                nextIdx1 = curIdx1 + 1
                nextIdx2 = curIdx2 + 1
    
                heapq.heappush(local.sum_idx, [nums1[nextIdx1] + nums2[nextIdx2], [nextIdx1, nextIdx2]])
                local.res.append([nextIdx1, nextIdx2])
    
            helper(k, curNum+1)
    
    
        helper(k, 0)
    
        # print "local.res = ", local.res
        #print "local.resVal = ", local.resVal
    
        return local.resVal

Log in to reply
 

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