P199 查找和最小的 K 对数字

描述#

这个题目说的是，给你两个递增的整数数组 nums1 和 nums2，以及一个整数 K。分别从 nums1 和 nums2 各取一个数字，可以构成一个数字对。你要找出 nums1 和 nums2 可以构成的所有数字对中，和最小的 K 对数字。

比如说，给你的两个数组和 K 分别是：
nums1 = [1, 5, 9]
nums2 = [2, 5, 8]
K = 4
这两个递增数组构成的所有数字对中，和最小的 4 对数字是：
[1, 2]
[1, 5]
[5, 2]
[1, 8]

相关视频：行列递增矩阵中第 K 小的元素

代码#

public class AlgoCasts {

  class Elem implements Comparable<Elem> {
    int idx1, idx2, sum;

    Elem(int idx1, int idx2, int sum) {
      this.idx1 = idx1;
      this.idx2 = idx2;
      this.sum = sum;
    }

    @Override
    public int compareTo(Elem o) {
      return this.sum - o.sum;
    }
  }

  // Time: O(n1*n2*log(k)), Space: O(k)
  public List<List<Integer>> kSmallestPairsMaxHeap(int[] nums1,
                                                   int[] nums2, int k) {
    if (nums1 == null || nums1.length == 0
      || nums2 == null || nums2.length == 0 || k <= 0) return new ArrayList<>();
    int n1 = nums1.length, n2 = nums2.length;
    Queue<Elem> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    for (int i = 0; i < n1; ++i) {
      for (int j = 0; j < n2; ++j) {
        int sum = nums1[i] + nums2[j];
        if (maxHeap.size() < k) {
          maxHeap.add(new Elem(i, j, sum));
        } else if (sum < maxHeap.peek().sum) {
          maxHeap.poll();
          maxHeap.add(new Elem(i, j, sum));
        }
      }
    }
    List<List<Integer>> result = new LinkedList<>();
    while (!maxHeap.isEmpty()) {
      Elem e = maxHeap.poll();
      result.add(0, Arrays.asList(nums1[e.idx1], nums2[e.idx2]));
    }
    return result;
  }

  // Time: O(k*log(k)), Space: O(k)
  public List<List<Integer>> kSmallestPairsMinHeap(int[] nums1,
                                                   int[] nums2, int k) {
    if (nums1 == null || nums1.length == 0
      || nums2 == null || nums2.length == 0 || k <= 0) return new ArrayList<>();
    int n1 = nums1.length, n2 = nums2.length;
    List<List<Integer>> result = new ArrayList<>();
    Queue<Elem> minHeap = new PriorityQueue<>();
    for (int i = 0; i < n1 && i < k; ++i)
      minHeap.add(new Elem(i, 0, nums1[i]+nums2[0]));

    for (int i = 0; i < k && !minHeap.isEmpty(); ++i) {
      Elem e = minHeap.poll();
      result.add(Arrays.asList(nums1[e.idx1], nums2[e.idx2]));
      ++e.idx2;
      if (e.idx2 < n2) {
        e.sum = nums1[e.idx1] + nums2[e.idx2];
        minHeap.add(e);
      }
    }
    return result;
  }

}