2018-08-31 13:47:59 +0000   |     algorithm leetcode array heap   |   Viewed times   |    

题目

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.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note:

计算所有可能的配对

老老实实计算出每一对和,放进一个Heap里输出。

代码

/**
 * x = min(m*n, k)
 * O(xlogx)
 * 老老实实计算出每一对和,放进一个Heap里输出。
 */
class Solution1 implements Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<>();
        if (nums1 == null || nums2 == null || k < 0) {
            return res;
        }
        PriorityQueue<int[]> heap = new PriorityQueue<>(new Comparator<int[]>(){
            public int compare(int[] a, int[] b) {
                return a[0] + a[1] - b[0] - b[1];
            }
        });
        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length; j++) {
                heap.offer(new int[]{nums1[i], nums2[j]});
            }
        }
        for (int i = 0; i < k && !heap.isEmpty(); i++) {
            res.add(heap.poll());
        }
        return res;
    }
}

结果

find-k-pairs-with-smallest-sums-1

Heap维护一组指针,O(klog(min(m,n)))

因为nums1nums2都是实现排过序的。所以,

nums1[x] + nums2[y] <= nums1[x] + nums2[y+1]

所以如果nums1[x] + nums2[y]都没被选上,就没必要拿nums1[x] + nums2[y+1]上去比。所以问题就抽象成一组从nums1到nums2的向量。比如,[1->2]这个向量代表的和被取走之后,向量右端(指向nums2)往上挪动一格,变成[1->9]。同理[7->2]被取走之后,更新为[7->9]。下一个最小的和一定出在这一组向量指向的组合中。 find-k-pairs-with-smallest-sums-a

再简化一点,甚至不需要从一开始就维护所有这些指针。一开始只需要一个指针,然后慢慢加入新指针。因为同理,

nums1[x] + nums2[0] <= nums1[x+1] + nums2[0]

代码

class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<>();
        if (nums1 == null || nums2 == null || nums1.length == 0 || nums2.length == 0 || k <= 0) {
            return res;
        }
        // 以较小的数组作为生成向量的主数组
        int[] numsA = (nums1.length <= nums2.length)? nums1 : nums2;
        int[] numsB = (numsA == nums1)? nums2 : nums1;
        PriorityQueue<int[]> heap = new PriorityQueue<>(new Comparator<int[]>(){
            public int compare(int[] a, int[] b) {
                return numsA[a[0]] + numsB[a[1]] - numsA[b[0]] - numsB[b[1]];
            }
        });
        int vectorP = 0;
        heap.offer(new int[]{vectorP, 0});
        while (k-- > 0 && !heap.isEmpty()) {
            int[] next = heap.poll();
            if (numsA == nums1) {
                res.add(new int[]{numsA[next[0]], numsB[next[1]]});
            } else {
                res.add(new int[]{numsB[next[1]], numsA[next[0]]});
            }
            if (next[0] == vectorP && vectorP < numsA.length - 1 && next[1] == 0) { // 添加新向量
                heap.offer(new int[]{++vectorP, 0});
            }
            if (next[1] < numsB.length - 1) { // 弹出的向量向前进一格
                heap.offer(new int[]{next[0], next[1] + 1});
            }
        }
        return res;
    }
}

结果

find-k-pairs-with-smallest-sums-2

一种可能的O(k)的解法

Stephanpochmann又给出了一种O(k)的解法。基本思想基于kth smallest element in a sorted matrix这个子问题,

整个和空间是一个二维矩阵。我先找到第k小的和,然后再从这个数一个个往前推。

原文链接 –> 【O(k) solution】

关键点就是他找到了一个在O(#row)时间里解决kth smallest element in a sorted matrix问题的办法。

原文链接 –> 【O(n) from paper. Yes, O(#rows).】