matrix - 如何在排序的 MxN 矩阵中找到第 K 个最小的和

Question

我看过有关如何在排序矩阵中找到第 K 个最小元素的解决方案，我还看过有关如何在两个数组中找到第 K 个最小和的解决方案。

但是我最近发现一个问题，要求在排序的 MxN 矩阵中找到第 K 个最小的和。总和必须由每一行中的一个元素组成。我真的在努力开发任何接近工作解决方案的东西，更不用说蛮力解决方案了。任何帮助将不胜感激!

我认为这会是某种堆问题...但也许这是一个图形问题？我不太擅长图表。

Answer 1

我假设“排序的 MxN 矩阵”是指矩阵的每一行都已排序。如果您已经知道如何合并 2 行并且只取前 K 个元素，则可以执行相同的过程来合并矩阵的每一行。忽略 int[] 和 List 之间的 Java 转换，下面的代码应该可以工作。

class Solution {

/**
 * Runtime O(m * k * logk)
 */ 
public int kthSmallest(int[][] mat, int k) {
    List<Integer> row = IntStream.of(mat[0]).boxed().collect(Collectors.toList());
    for (int i = 1; i < mat.length; i++) {
        row = kthSmallestPairs(row, mat[i], k);
    }
    return row.get(k - 1);
}

/**
 * A pair is formed from one num of n1 and one num of n2. Find the k-th smallest sum of these pairs
 * Queue size is maxed at k, hence this method run O(k logk)
 */
List<Integer> kthSmallestPairs(List<Integer> n1, int[] n2, int k) {
    // 0 is n1's num,     1 is n2's num,     2 is n2's index
    Queue<int[]> que = new PriorityQueue<>((a, b) -> a[0] + a[1] - b[0] - b[1]);

    // first pair each num in n1 with the 0-th num of n2. Don't need to do more than k elements because those greater
    // elements will never have a chance
    for (int i = 0; i < n1.size() && i < k; i++) {
        que.add(new int[] {n1.get(i), n2[0], 0});
    }

    List<Integer> res = new ArrayList<>();
    while (!que.isEmpty() && k-- > 0) {
        int[] top = que.remove();
        res.add(top[0] + top[1]);
        // index of n2 is top[2]
        if (top[2] < n2.length - 1) {
            int nextN2Idx = top[2] + 1;
            que.add(new int[] {top[0], n2[nextN2Idx], nextN2Idx});
        }
    }

    return res;
}