Find K Pairs with Smallest Sums

最无脑的写法，复杂度 O(k * log(mn))，minHeap 里把所有可能的 pair 都存进去了。

我自己用heap加了优化, 整体思路是对的, 错了一个细节. 我开始是push 1个pair element<a0, b0, val>进去, 然后每次移动a0, b0. 这样出的问题是 a0b0->a0b1, a1b0->a0b2, a1b1, a2b0, a1b1. a1b1算了两次.

public class Solution {

    class Entity{
        int idx1;
        int idx2;
        int val;
        public Entity(int idx1, int idx2, int val){
            this.idx1 = idx1;
            this.idx2 = idx2;
            this.val = val;
        }
    }

    private class MyComparator implements Comparator<Entity>{
        public int compare(Entity a, Entity b){
            return a.val - b.val;
        }
    }

    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> ans = new ArrayList<int[]>();
        if(nums1 == null || nums2 == null || 
           nums1.length == 0 || nums2.length == 0) 
           return ans;
        PriorityQueue<Entity>pq = new PriorityQueue<Entity>(new MyComparator());

        pq.offer(new Entity(0, 0, nums1[0]+nums2[0]));
        while(k>0&&!pq.isEmpty()){
            Entity e = pq.poll();

            if(e.idx2<nums2.length-1)
                pq.offer(new Entity(e.idx1, e.idx2+1, nums1[e.idx1]+nums2[e.idx2+1]));
            if(e.idx1<nums1.length-1)
                pq.offer(new Entity(e.idx1+1, e.idx2, nums1[e.idx1+1]+nums2[e.idx2]));
            ans.add(new int[]{nums1[e.idx1], nums2[e.idx2]});
            k--;
        }
        return ans;
    }
}

稍作修改, 开始offer k个element进heap, <a0, b0>....<ak, b0>. 然后每次b0++, 轻松AC 那么显然的优化是，既然我们只要 k 个答案，heap 里存 k 个就行，同时要尽可能的利用数组已经排好序的特性， 让 index 单调向前走。

O(k log k) ~ 超过 81%

自定义 Tuple，存 int[] pair 还有当前 pair 相对于 nums2 的 index

• 先以 nums1 的前 k 个数和 nums2 的第一个数为起点，初始化 minHeap；

• 每次新 tuple 出来的时候，都把 index2 往后移一步，index1 不变 (还取 tuple 里 pair[0] 的值)

• 在这个过程中，一个 Pair 的 ptr1 和 ptr2 是单调递增的，因为 heap 里已经存了之前看过的组合了，不能回头，否则会有重复答案。

public class Solution {

    class Entity{
        int idx1;
        int idx2;
        int val;
        public Entity(int idx1, int idx2, int val){
            this.idx1 = idx1;
            this.idx2 = idx2;
            this.val = val;
        }
    }

    private class MyComparator implements Comparator<Entity>{
        public int compare(Entity a, Entity b){
            return a.val - b.val;
        }
    }

    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> ans = new ArrayList<int[]>();
        if(nums1 == null || nums2 == null || 
           nums1.length == 0 || nums2.length == 0) 
           return ans;
        PriorityQueue<Entity>pq = new PriorityQueue<Entity>(new MyComparator());
        for(int i=0;i<nums1.length&&i<k;i++)
            pq.offer(new Entity(i, 0, nums1[i]+nums2[0]));

        while(k>0&&!pq.isEmpty()){
            Entity e = pq.poll();

            if(e.idx2<nums2.length-1)
                pq.offer(new Entity(e.idx1, e.idx2+1, nums1[e.idx1]+nums2[e.idx2+1]));

            ans.add(new int[]{nums1[e.idx1], nums2[e.idx2]});
            k--;
        }
        return ans;
    }
}

Kth Smallest Element in a Sorted Matrix

public class Solution {
    private class Tuple implements Comparable<Tuple>{
        int x, y, val;

        public Tuple(int x, int y, int val){
            this.val = val;
            this.x = x;
            this.y = y;
        }
        public int compareTo(Tuple tuple){
            return this.val - tuple.val;
        }
    }

    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Tuple> minHeap = new PriorityQueue<Tuple>();
        int rows = matrix.length;
        for(int i = 0; i < rows; i++) minHeap.offer(new Tuple(i, 0, matrix[i][0]));
        for(int i = 0; i < k - 1; i++){
            Tuple tuple = minHeap.poll();
            if(tuple.y + 1 == matrix[0].length) continue;
            minHeap.offer(new Tuple(tuple.x, tuple.y + 1, matrix[tuple.x][tuple.y + 1]));
        }

        return minHeap.peek().val;
    }


}

Kth Smallest Element in a Sorted Matrix

同一道题，扩展到 k 个 array 的情况~

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix[0].length;
        PriorityQueue<Tuple> pq = new PriorityQueue<Tuple>();
        for(int j = 0; j <= n-1; j++) pq.offer(new Tuple(0, j, matrix[0][j]));
        for(int i = 0; i < k-1; i++) {
            Tuple t = pq.poll();
            if(t.x == n-1) continue;
            pq.offer(new Tuple(t.x+1, t.y, matrix[t.x+1][t.y]));
        }
        return pq.poll().val;
    }
}

class Tuple implements Comparable<Tuple> {
    int x, y, val;
    public Tuple (int x, int y, int val) {
        this.x = x;
        this.y = y;
        this.val = val;
    }

    @Override
    public int compareTo (Tuple that) {
        return this.val - that.val;
    }
}