在数组的所有可能排序对的有序列表中找到第K个对

给定一个包含n个整数的数组nums，我们可以将其分为(i, j)的多个有序对，其中i < j，并且nums[i] < nums[j]。

请编写一个函数，找到在数组的所有可能排序对的有序列表中排名第k的有序对。

算法思路

一种比较简单的思路是暴力枚举所有的有序对，然后按照它们的大小排序，找到第k个有序对就可以了。时间复杂度为O(n^2logn)，显然不够优秀。

更好的思路是利用归并排序的思想，在归并排序过程中进行计算。首先将数组分成左右两个部分进行排序，然后将左右两个部分合并时，计算出左右两个部分之间的有序对数。

具体来说，将数组分成左右两个部分，分别进行排序，并计算出左右两个部分之间的有序对数。然后将两个部分合并成一个有序数组，同时计算出左右两个部分之间的新有序对数。重复进行上述过程，直到数组长度为1为止。最后得到的有序数组中，第k个有序对就是我们要找的结果。

时间复杂度为O(nlogn)，相比暴力枚举的方法要快很多。

代码实现

下面是Python的实现，代码中的merge函数是归并排序中的合并函数，count函数用于计算有序对数。由于Python中函数默认为全局变量，代码本身比较简单，易于理解。

class Solution:
    def kthSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        if not nums1 or not nums2:
            return []
        
        def count(nums1, nums2, target):
            i, cnt = len(nums1) - 1, 0
            for j in range(len(nums2)):
                while i >= 0 and nums1[i] + nums2[j] > target:
                    i -= 1
                cnt += i + 1
            return cnt
            
        def merge(nums1, nums2, target):
            ans, i, j = [], 0, 0
            while i < len(nums1) and j < len(nums2):
                if nums1[i] + nums2[j] <= target:
                    ans.append([nums1[i], nums2[j]])
                    i += 1
                else:
                    j += 1
            return ans
            
        left, right = 2 * nums1[0], 2 * nums1[-1] + 1
        mid = (left + right) // 2
        while left < right:
            cnt = count(nums1, nums2, mid)
            if cnt < k:
                left = mid + 1
            else:
                right = mid
            mid = (left + right) // 2
            
        ans = merge(nums1, nums2, mid)
        n = len(ans)
        if n < k:
            ans += [[nums1[i], nums2[-1]] for i in range(min(k - n, len(nums1)))]
            ans += [[nums1[-1], nums2[i]] for i in range(min(k - n, len(nums2)))]
        return ans[:k]

参考资料

LeetCode 373. Find K Pairs with Smallest Sums