Heap，排序 matrix 中的 two pointers

Find K Pairs with Smallest Sums

这题一开始搞错了，还以为是 two pointer 问题，其实同一个元素是可以复用的，只要两个 index 不一样就行了。因此总共 pair 的数量是 m * n，光是靠单向移动的 two pointer 是不行的，会漏掉正解。得靠 heap.

Given nums1 = [1,1,2], nums2 = [1,2,3], k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

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

public class Solution {
    private class MyComparator implements Comparator<int[]>{
        public int compare(int[] a, int[] b){
            return a[0] + a[1] - b[0] - b[1];
        }
    }
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> list = new ArrayList<>();

        PriorityQueue<int[]> minHeap = new PriorityQueue<int[]>(new MyComparator());

        for(int i = 0; i < nums1.length; i++){
            for(int j = 0; j < nums2.length; j++){
                minHeap.offer(new int[]{nums1[i], nums2[j]});
            }
        }

        while(!minHeap.isEmpty() && k > 0){
            list.add(minHeap.poll());
            k--;
        }

        return list;
    }
}

那么显然的优化是，既然我们只要 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 里已经存了之前看过的组合了，不能回头，否则会有重复答案。

这是一个借用 heap 的 two pointer 算法。

public class Solution {
    private class Tuple{
        int val;
        int ptr1;
        int ptr2;
        public Tuple(int val, int ptr1, int ptr2){
            this.val = val;
            this.ptr1 = ptr1;
            this.ptr2 = ptr2;
        }
    }

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

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

        PriorityQueue<Tuple> minHeap = new PriorityQueue<Tuple>(new MyComparator());

        for(int i = 0; i < nums1.length && i < k; i++){
                minHeap.offer(new Tuple(nums1[i] + nums2[0], i, 0));
        }

        while(!minHeap.isEmpty() && k > 0){
            Tuple tuple = minHeap.poll();
            int ptr1 = tuple.ptr1;
            int ptr2 = tuple.ptr2;
            list.add(new int[]{nums1[ptr1], nums2[ptr2]});

            if(ptr2 + 1 < nums2.length){
                minHeap.offer(new Tuple(nums1[ptr1] + nums2[ptr2 + 1], ptr1, ptr2 + 1));
            }
            k--;
        }

        return list;
    }
}

Kth Smallest Element in a Sorted Matrix

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

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;
    }


}