8.2.3 python排序LeetCode题目(3) —— K Pairs with Smallest Sums & Kth Smallest Element in a Sorted Matrix

这一节再看一些排序相关的题目。

 

373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: The first 3 pairs are returned from the sequence: 
             [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

题目解析：

我们先看用堆排序的一个简单方法。原理比较简单，就是遍历所有pairs，最后使用 heapq的nsmallest方法。另外当k较小时，做了一些简化，不必遍历所有数。在此基础上进行优化，看下面方法二，仍然要遍历所有的pairs，维持res的长度为k，当res满了时，和res里最大的比较（因此数值取负数，因为要一个最大堆），新的数更小时，pop出最大的push新的值。方法二性能显著提升，但是仍然距离最好的方法有差距。

# 方法一
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        import heapq
        m, n = len(nums1), len(nums2)
        res = []
        
        if k <= m and k <= n:
            pairs = []
            fl = k
            for num1 in nums1[:k]:
                for num2 in nums2[:fl]:
                    pairs.append([num1, num2])
                fl -= 1
            
        else:
            pairs = []
            for num1 in nums1:
                for num2 in nums2:
                    pairs.append([num1, num2])
                    
        return heapq.nsmallest(k, pairs, key=lambda x: x[0]+x[1])

# 方法二
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        N, M = len(nums1), len(nums2)
        n, m = min(N, k), min(M, k)
        res = []
        for i in range(n):
            for j in range(m):
                if len(res) < k:
                    heapq.heappush(res, (-(nums1[i]+nums2[j]), i, j))
                else:
                    if nums1[i] + nums2[j] < -res[0][0]:
                        heapq.heappop(res)
                        heapq.heappush(res, (-(nums1[i]+nums2[j]), i, j))
        ans = []
        while len(res):
            _, i, j = heapq.heappop(res)
            ans.append([nums1[i], nums2[j]])
        return ans[::-1]

 接下来我们看方法三，耗时较少的方算法基本都是这一思路。nums1[0]+nums2[0] 一定是最小的，那么前k个最小的，是nums1和nums2中哪些数相加的呢？前两张方法中，都没有用到两个数组已排序的条件，因此我们要动态的根据两个数组中数值的大小关系找前k个最小值。

每次循环从que中弹出当前最小的值，根据该pair的i，j的索引，我们尝试(i, j+1) ，若j一直是0， 则还要尝试 (i+1,j)。在这两个pair及que中仍存在的pair中排出下一个最小的。这样不必遍历所有数，耗时相当少。

# 方法三
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        if not (nums1 and nums2): return []
        n, m, res, que = len(nums1), len(nums2), [], [(nums1[0]+nums2[0], 0, 0)]
        for _ in range(min(k, n*m)):
            _, i, j = heapq.heappop(que)
            res.append([nums1[i], nums2[j]])
            if j+1 < m:
                heapq.heappush(que, (nums1[i]+nums2[j+1], i, j+1))
            if j == 0 and i+1 < n:
                heapq.heappush(que, (nums1[i+1]+nums2[j], i+1, j))
        return res

 

378. Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

题目解析：

kth，这类题一般用堆来解决。和上面的题类似，方法一，我们暴力遍历矩阵，奇怪的是，性能还不错。

# 方法1
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        import heapq
        data = []
        for row in matrix:
            for num in row:
                heapq.heappush(data, num)
        i = 0
        while i < k:
            n = heapq.heappop(data)            
            i += 1
        return n

再看方法二，首先请参考 6.2.3 python二分查找算法及LeetCode题目(3)之二维数组 —— Search a 2D Matrix 。

在这样的二维矩阵中，查找一个数，时间复杂度O(m+n)，从第一行最后一个元素开始查找，大于target光标左移，小于target光标下移。

因此方法二中，我们取矩阵最小值和最大值平均，不断二分查找，看小于该数的数在矩阵中的个数，对应于抽出的函数 counterSmaller，直到我们找到结果。该方法让我们进一步理解二维矩阵的查找问题。。

# 方法2
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        def countSmaller(num):
            i, j, cnt = 0, n-1, 0
            while i<n and j>=0:
                if matrix[i][j]<num:
                    cnt+=j+1
                    i+=1
                else:
                    j-=1
            return cnt
        
        low, high=matrix[0][0], matrix[-1][-1]
        while low <= high:
            mid = (low + high) // 2
            if countSmaller(mid)>=k:
                high = mid - 1
            else:
                low = mid + 1
        return high

---

这一节两道题与堆有关，但是可以不用堆；数组是已排序的，不需要我们排序，但是我们解题要充分利用有序的条件，this is most important!

写在第八章的最后：

排序到此结束，题目熟悉即可，关键是几大排序算法的实现及性能分析吧，比如直接让你写个快排的函数，这和前几章的要求有所不同。学习到此，基础的数据结构与算法基本结束，也不必追求更多的数据结构，如图论，也不必追求更多的算法，虽然后面的动态规划还是值得学一学的，把前面的知识和经典题目牢记更重要哈。后续文章可能会对前面知识点及经典题目进行总结，帮助复习。