【每日一题Day220】LC1439有序矩阵中的第 k 个最小数组和 | 堆

再来做一下373，之前都没有试过用小顶堆求第K小的，有序这个条件对我而言是摆设了

查找和最小的 K 对数字【LC373】

给定两个以 升序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk) 。

大顶堆

• 思路：使用大顶堆存放最小的K对数字，堆将数对之和从大到小排序，堆顶为最大值，当堆的大小大于$k$并且当遍历到的数对之和小于堆顶数对之和时，将堆顶数对弹出，将新数对压入；当堆的大小小于$k$时，直接将新数对压入

  • 由于数组是升序排列，因此每个数组只需要取前$min(k,n)$个元素即可

• 实现

  每个循环最多执行k次

  class Solution {
  
      public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
          PriorityQueue<int[]> pq = new PriorityQueue<>((o1, o2) -> o2[0] + o2[1] - o1[0] -o1[1]);
          List<List<Integer>> res = new ArrayList<>();
          int n = nums1.length, m = nums2.length;
          for (int i = 0; i < Math.min(k, n); i++){
              for (int j = 0; j < Math.min(k, m); j++){
                  int[] n1 = new int[]{nums1[i], nums2[j]};
                  if (pq.size() < k){            
                      pq.add(n1);
                  }else if(pq.size() == k && n1[0] + n1[1] < pq.peek()[0] + pq.peek()[1]){
                      pq.poll();
                      pq.add(n1); 
                  }
              }
          }
          while(!pq.isEmpty()){
              int[] poll = pq.poll();
              res.add(Arrays.asList(poll[0], poll[1]));
          }
          return res;
      }
  }
  

  • 复杂度
    • 时间复杂度：$O(k^{2}logk)$，向堆中添加元素的时间复杂度为$O(logk)$
    • 空间复杂度：$O(k)$

小顶堆：归并

多路归并：先从一个数组中取最小的k个，然后叠加第二个数组

• 思路：

  • 由于数组是升序的，因此最小的数对一定是 { 0 , 0 } \{0,0\} {0,0}，那么第二小的数对可能是 { 0 , 1 } \{0,1\} {0,1}或者 { 1 , 0 } \{1,0\} {1,0}。
  • 因此可以使用小顶堆按从小到大的顺序快速求出较小的$k$个数对，小顶堆存储三元组$[值,nums1下标,nums2下标]$。
  • 当出堆时，将下一个可能的数对放入堆中， { i , j + 1 } \{i,j+1\} {i,j+1}或者 { i + 1 , j } \{i+1,j\} {i+1,j}【重要】
    • 当 { i , j } \{i,j\} {i,j}入堆时，出堆的下标可能为 { i , j − 1 } \{i,j-1\} {i,j−1}或者 { i − 1 , j } \{i-1,j\} {i−1,j}
    • 为了避免重复放入，保证 { i , j − 1 } \{i,j-1\} {i,j−1}和 { i − 1 , j } \{i-1,j\} {i−1,j}出堆时，只有一个会将 { i , j } \{i,j\} {i,j}入堆【可以使用哈希表记录在堆中的下标或者进行特殊规定】
    • 规定 { i , j − 1 } \{i,j-1\} {i,j−1}入堆时，将 { i , j } \{i,j\} {i,j}入堆。初始化时我们需要将 { 0 , 0 } \{0,0\} {0,0}， { 1 , 0 } \{1,0\} {1,0}， { 2 , 0 } \{2,0\} {2,0}……进行入堆【该初始化相当于先在第一个数组中选择较小的$k$个元素】
    • 然后出堆时，将该数对记录在结果集中，该数对为最小的第$i$个数对。按照规则进行入堆【相当于组合第二个数组中的元素】

• 实现

  class Solution {
      public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
          int n = nums1.length, m = nums2.length;
          var ans = new ArrayList<List<Integer>>(k); // 预分配空间
          var pq = new PriorityQueue<int[]>((a, b) -> a[0] - b[0]);
          for (int i = 0; i < Math.min(n, k); i++) // 至多 k 个
              pq.add(new int[]{nums1[i] + nums2[0], i, 0});
          while (!pq.isEmpty() && ans.size() < k) {
              var p = pq.poll();
              int i = p[1], j = p[2];
              ans.add(List.of(nums1[i], nums2[j]));
              if (j + 1 < m)
                  pq.add(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
          }
          return ans;
      }
  }
  
  作者：灵茶山艾府
  链接：https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/solutions/2286318/jiang-qing-chu-wei-shi-yao-yi-kai-shi-ya-i0dj/
  来源：力扣（LeetCode）
  著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
  

  • 复杂度
    • 时间复杂度：$\mathcal{O}(klogmin(n,k))$，为了获得$k$个数对，需要循环$k$次每次出堆入堆的的时间复杂度为$\mathcal{O}(logmin(n,k))$
    • 空间复杂度：$\mathcal{O}(min(n,k))$

• 优化

  在循环的过程中将$(i,0)$入堆。

  class Solution {
      public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
          int n = nums1.length, m = nums2.length;
          var ans = new ArrayList<List<Integer>>(k); // 预分配空间
          var pq = new PriorityQueue<int[]>((a, b) -> a[0] - b[0]);
          pq.add(new int[]{nums1[0] + nums2[0], 0, 0});
          while (!pq.isEmpty() && ans.size() < k) {
              var p = pq.poll();
              int i = p[1], j = p[2];
              ans.add(List.of(nums1[i], nums2[j]));
              if (j == 0 && i + 1 < n)
                  pq.add(new int[]{nums1[i + 1] + nums2[0], i + 1, 0});
              if (j + 1 < m)
                  pq.add(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
          }
          return ans;
      }
  }
  
  作者：灵茶山艾府
  链接：https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/solutions/2286318/jiang-qing-chu-wei-shi-yao-yi-kai-shi-ya-i0dj/
  来源：力扣（LeetCode）
  著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
  

有序矩阵中的第 k 个最小数组和【LC1439】

给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。

你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。

• 思路
  • 我们需要从m行中每行选取1个元素，每行有n个选择，因此总共方案数为$n^m$，由于需要找到的是第 k 个最小数组和，因此只需要记录最小的$k$个和即可。
  • 只需要记录选到第$i$行时，最小的$k$个和即可（比这些元素大的一定不在结果中），答案一定是在这些组合的基础上添加后续元素，然后搜索下一行（进行多路归并），从小到大记录最小的$k$个和，最后返回第$k$个和。

暴力

• 实现：暴力

  class Solution {
      public int kthSmallest(int[][] mat, int k) {
          // 从m行中每行选取1个元素，每行有n个选择，因此总共方案数为n^m
          // 我们最终要返回方案和为第k小的和，因此只需要记录选到第i行时，最小的k个和即可，然后进行多路归并，结果一定是在这些组合的基础上添加元素
          // 实现：暴力或者小顶堆
          int m = mat.length, n = mat[0].length;
          int[] a = new int[]{0};
          for (int[] row : mat){
              int[] b = new int[a.length * n];
              int i = 0;
              for (int num1 : a){
                  for (int num2 : row){
                      b[i++] = num1 + num2;
                  }
              }
              Arrays.sort(b);
              if (b.length > k){
                  b = Arrays.copyOfRange(b, 0, k);
              }
              a = b;
              
          } 
          return a[k - 1];
      }
  }
  

  • 复杂度
    • 时间复杂度：$\mathcal{O}(m\ast min(n,k)\ast k\ast log(min(n,k)\ast k))$。m和n分别为矩阵的行和列。遍历前m行中$min(n,k)$个元素，与前一次的结果进行组合，因此时间复杂度为$\mathcal{O}(m\ast min(n,k)\ast k)$。排序的时间复杂度为$\mathcal{O}(log(min(n,k)\ast k))$
    • 空间复杂度：$\mathcal{O}(min(n,k)\ast k)$

小顶堆：归并

• 实现：小顶堆

  使用小顶堆进行优化，具体思路同[查找和最小的 K 对数字【LC373】](#查找和最小的 K 对数字【LC373】)，相当于将LC373循环m次

  • 使用小顶堆按从小到大的顺序快速求出较小的$k$个和，小顶堆存储三元组$[值,nums1下标,nums2下标]$。
    • 出栈的元素，即为遍历到第$i$行时第k个小的和
    • 出栈的同时，将下一行的元素加入

  class Solution {
      public int kthSmallest(int[][] mat, int k) {
          var a = new int[]{0};
          for (var row : mat)
              a = kSmallestPairs(row, a, k);
          return a[k - 1];
      }
  
      // 373. 查找和最小的 K 对数字
      private int[] kSmallestPairs(int[] nums1, int[] nums2, int k) {
          int n = nums1.length, m = nums2.length, sz = 0;
          var ans = new int[Math.min(k, n * m)];
          var pq = new PriorityQueue<int[]>((a, b) -> a[0] - b[0]);
          pq.add(new int[]{nums1[0] + nums2[0], 0, 0});
          while (!pq.isEmpty() && sz < k) {
              var p = pq.poll();
              int i = p[1], j = p[2];
              ans[sz++] = nums1[i] + nums2[j]; // 数对和
              if (j == 0 && i + 1 < n)
                  pq.add(new int[]{nums1[i + 1] + nums2[0], i + 1, 0});
              if (j + 1 < m)
                  pq.add(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
          }
          return ans;
      }
  }
  
  作者：灵茶山艾府
  链接：https://leetcode.cn/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solutions/2286593/san-chong-suan-fa-bao-li-er-fen-da-an-du-k1vd/
  来源：力扣（LeetCode）
  著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
  

  • 复杂度
    • 时间复杂度：$\mathcal{O}(m\ast k\ast log(min(n,k)))$，为了获得$k$个数对，需要循环$k$次每次出堆入堆的的时间复杂度为$\mathcal{O}(logmin(n,k))$
    • 空间复杂度：$\mathcal{O}(min(n,k))$