leetcode--有序矩阵中的第 k 个最小数组和

 题目是LeetCode第187场周赛的第四题，链接：有序矩阵中的第 k 个最小数组和。具体描述为：给你一个m * n的矩阵 mat，以及一个整数k，矩阵中的每一行都以非递减的顺序排列。你可以从每一行中选出1个元素形成一个数组。返回所有可能数组中的第k个最小数组和。

• m == mat.length
• n == mat.length[i]
• 1 <= m, n <= 40
• 1 <= k <= min(200, n ^ m)
• 1 <= mat[i][j] <= 5000
• mat[i] 是一个非递减数组

 示例1:

输入：mat = [[1,3,11],[2,4,6]], k = 5
输出：7
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。  

 示例2:

输入：mat = [[1,3,11],[2,4,6]], k = 9
输出：17

 示例3:

输入：mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出：9
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。 

 示例4:

输入：mat = [[1,1,10],[2,2,9]], k = 7
输出：12

 这道题因为有k不超过200的限制所以可以直接暴力解决，维护一个列表保存k个最小和，每次遍历一行的时候遍历到的n个数都跟这k个最小和相加得到一共nk个最小和，然后排序，截取前k个最小和，继续下去直到最后一行。时间复杂度为$O(mnklog(nk))$，空间复杂度为$O(nk)$。

 JAVA版代码如下：

class Solution {
    public int kthSmallest(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        List<Integer> list = new ArrayList<>();
        list.add(0);
        for (int i = 0; i < m; ++i) {
            List<Integer> temp = new ArrayList<>();
            for (int j = 0; j < n; ++j) {
                for (int item : list) {
                    temp.add(item + mat[i][j]);
                }
            }
            Collections.sort(temp);
            list = temp.subList(0, Math.min(k, temp.size()));
        }
        return list.get(k - 1);
    }
}

 提交结果如下:

 然后也可以用一个小顶堆来维护最小和，这样用k次循环，每次循环将堆顶的出堆得到当前最小和，最后一次出堆的就是我们需要的。然后在每次循环里，需要得到出堆最小和对应的索引序列（比如一开始就是[0,0,...,0]），将这个索引序列各个位置分别+1后加入堆里去（同时需要去重，所以每次最多会有m个新的加入堆去）。时间复杂度为$O(mklog(mk))$，空间复杂度为$O(m^{2}k)$。

 JAVA版代码如下：

class Solution {
    private String toStr(int[] idx) {
        StringBuilder sb = new StringBuilder();
        for (int i : idx) {
            sb.append(String.valueOf(i) + "-");
        }
        return sb.toString();
    }
    public int kthSmallest(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        // 小顶堆
        PriorityQueue<SumAndIndex> heap = new PriorityQueue<>(
            new Comparator<SumAndIndex> () {
                @Override
                public int compare(SumAndIndex s1, SumAndIndex s2) {
                    return s1.curSum - s2.curSum;
                }
            }
        );
        int curSum = 0;
        for (int i = 0; i < m; ++i) {
            curSum += mat[i][0];
        }
        int[] index = new int[m];
        SumAndIndex item = new SumAndIndex(curSum, index);
        heap.offer(item);
        Set<String> set = new HashSet<>();
        set.add(toStr(index));
        for (int count = 0; count < k - 1; ++count) {
            item = heap.poll();
            index = item.indexs;
            for (int i = 0; i < m; ++i) {
                if (index[i] < n - 1) {
                    int[] newIdx = index.clone();
                    ++newIdx[i];
                    if (set.add(toStr(newIdx))) {
                        int newSum = item.curSum - mat[i][index[i]] + mat[i][newIdx[i]];
                        SumAndIndex newItem = new SumAndIndex(newSum, newIdx);
                        heap.offer(newItem);
                    }
                }
            }
        }
        item = heap.poll();
        return item.curSum;
    }
}

class SumAndIndex {
    int curSum;
    int[] indexs;
    public SumAndIndex(int s, int[] i) {
        curSum = s;
        indexs = i;
    }
}

 提交结果如下:

 最后是二分法，先计算可能的最小值和最大值（也就是第一列之和以及最后一列之和），在这两个值之间进行二分，为了找到第k小值，我们需要计算小于某个值mid的所有可能组合有多少，这就是dfs()函数做的事情了。其做法是先遍历本行，只要和不超过mid的就可以进入下一行继续搜索，一旦发现有至少k个小于mid的需要及时剪枝。时间复杂度为$O(klog(N))$（其中N为和的最大可能值，这里是$40\times 5000$），空间复杂度为$O(k)$，因为每次dfs最多只会搜索k次。

 JAVA版代码如下：

class Solution {
    private int m;
    private int n;
    private int count;

    private void dfs(int[][] mat, int row, int k, int sum, int mid) {
        if (row == m || count >= k || sum > mid) {
            return;
        }
        dfs(mat, row + 1, k, sum, mid);
        for (int j = 1; j < n; ++j) {
            if (sum + mat[row][j] - mat[row][0] <= mid) {
                ++count;
                dfs(mat, row + 1, k, sum + mat[row][j] - mat[row][0], mid);
            }
            else {
                break;
            }
        }
    }

    public int kthSmallest(int[][] mat, int k) {
        m = mat.length;
        n = mat[0].length;
        int left = 0, right = 0;
        for (int i = 0; i < m; ++i) {
            left += mat[i][0];
            right += mat[i][n - 1];
        }
        int minSum = left;
        while (left < right) {
            int middle = left + (right - left) / 2;
            count = 1;
            dfs(mat, 0, k, minSum, middle);
            if (count < k) {
                left = middle + 1;
            }
            else {
                right = middle;
            }
        }
        return right;
    }
}

 提交结果如下:

 Python版代码如下：

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        self.m = len(mat)
        self.n = len(mat[0])
        self.count = 0
        self.k = k

        def dfs(row, curSum, mid):
            if row == self.m or curSum > mid or self.count >= self.k:
                return
            dfs(row + 1, curSum, mid)
            for j in range(1, self.n):
                if curSum + mat[row][j] - mat[row][0] <= mid:
                    self.count += 1
                    dfs(row + 1, curSum + mat[row][j] - mat[row][0], mid)
                else:
                    break

        left = sum(mat[i][0] for i in range(self.m))
        right = sum(mat[i][-1] for i in range(self.m))
        minSum = left
        while left < right:
            mid = (left + right) // 2
            self.count = 1
            dfs(0, minSum, mid)
            if self.count < k:
                left = mid + 1
            else:
                right = mid
        return right

 提交结果如下: