Leetcode 373. Find K Pairs with Smallest Sums（Python）

最新推荐文章于 2024-08-25 06:41:42 发布

wds2006sdo

最新推荐文章于 2024-08-25 06:41:42 发布

阅读量3.3k

点赞数

分类专栏： leetcode 文章标签： leetcode heap

本文链接：https://blog.csdn.net/wds2006sdo/article/details/51951753

版权

leetcode 专栏收录该内容

6 篇文章 0 订阅

订阅专栏

睡不着咋办！

题目大意

给定两个数组，定义pair是一个二元组，每个元素分别取自两个数组。求pair和前k小的pair

思路

用堆，我们维持一个k最大堆即可

代码

class mHeap(object):
    heap = []

    def __init__(self,MIN_HEAP = True):
        self.heap = []
        self.MIN_HEAP = MIN_HEAP

    def cal(self,i,j):
        if self.MIN_HEAP:
            return i[0] < j[0]
        return i[0] > j[0]

    def cal_equal(self,i,j):
        if self.MIN_HEAP:
            return i[0] <= j[0]
        return i[0] >= j[0]

    def top(self):
        return self.heap[0]

    def push(self,element):
        i = len(self.heap)
        self.heap.append(element)

        j = (i-1)/2
        tmp = element
        while( j >= 0 and i != 0 ):
            if self.cal(self.heap[j],tmp):
                break
            self.heap[i] = self.heap[j]
            i =j
            j = ( i - 1 ) / 2
        self.heap[i] = tmp

    def pop(self):
        if len(self.heap) == 0:
            return

        self.heap[0] = self.heap[-1]
        self.heap.pop()
        if len(self.heap) == 0:
            return
        self.fix_down()

    def fix_down(self):
        i = 0
        j = 2 * i + 1
        n = len(self.heap)

        tmp = self.heap[i]

        while j < n:
            if j + 1 < n and self.cal(self.heap[j+1],self.heap[j]):
                j += 1
            if self.cal_equal(tmp,self.heap[j]):
                break

            self.heap[i] = self.heap[j]
            i = j
            j = 2 * i + 1

        self.heap[i] = tmp

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        heap = mHeap(False)

        for i in nums1:
            for j in nums2:
                if len(heap.heap) < k:
                    heap.push([i+j,[i,j]])
                else:
                    top = heap.top()
                    if i+j < top[0]:
                        heap.pop()
                        heap.push([i+j,[i,j]])

        result = []
        for ele in heap.heap:
            result.append(ele[1])
        return result