【算法】Find the Kth Smallest Sum of a Matrix With Sorted Rows 有序矩阵中的第 k 个最小数组和

Find the Kth Smallest Sum of a Matrix With Sorted Rows 有序矩阵中的第 k 个最小数组和

问题描述：

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

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

m ， n 范围[1,40]

k 范围[1, min(200,n^m)]

分析

首先矩阵的行列范围是$[1,40]$,而且矩阵的每一行是一个有序 排列。

要求每一行选1个元素组成一个数组，而且这个数组的和是所有排列中第k个最小的。

简单暴力的方式就是，穷举所有排列，然后排序，找到第k个最小的。

记$行数为m，列数为n$，那么一个排列中会有m个元素，排列总数就是$N\ast N\ast N..\ast N$,m个n相乘，以M,N=40，总数就是$(40{)}^{40}$，所以这个方法走不通。

暴力走不通，那就只能找规律，这里可以借鉴大小堆的思路。

对于长度为$n$的排列A，如果A是第k个最小的排列，那么A中前$n-1$的元素和，在$n-1$个元素和中一定不会超过第k个。

因此，每个长度为$n^{\prime }$的排列对应的元素和，只需要保留前k个最小的，那么再加上1个元素，所产生的新排列，规模就是$k\ast n^{\prime }$,以给出的数据规模$k\ast n^{\prime }$上限就是$8000$，一共计算$m轮，最大规模上限8000\ast 40$.

代码

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

代码来源于 灵神大佬
这里的实现没有使用哈希表，你细品。
时间复杂度 O($M\ast kLog(min(k,N))$)

空间复杂度： O($min(k,N)$)

Tag

Array matrix heap