Question

1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

中文文档

Description

You are given an m x n matrix mat that has its rows sorted in non-decreasing order and an integer k.

You are allowed to choose exactly one element from each row to form an array.

Return _the _kth_ smallest array sum among all possible arrays_.

 

Example 1:

Input: mat = [[1,3,11],[2,4,6]], k = 5
Output: 7
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,2], [1,4], [3,2], [3,4], [1,6]. Where the 5th sum is 7.

Example 2:

Input: mat = [[1,3,11],[2,4,6]], k = 9
Output: 17

Example 3:

Input: mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
Output: 9
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]. Where the 7th sum is 9.  

 

Constraints:

• m == mat.length
• n == mat.length[i]
• 1 <= m, n <= 40
• 1 <= mat[i][j] <= 5000
• 1 <= k <= min(200, nm)
• mat[i] is a non-decreasing array.

Solutions

Solution 1

class Solution:
  def kthSmallest(self, mat: List[List[int]], k: int) -> int:
    pre = [0]
    for cur in mat:
      pre = sorted(a + b for a in pre for b in cur[:k])[:k]
    return pre[-1]

class Solution {
  public int kthSmallest(int[][] mat, int k) {
    int m = mat.length, n = mat[0].length;
    List<Integer> pre = new ArrayList<>(k);
    List<Integer> cur = new ArrayList<>(n * k);
    pre.add(0);
    for (int[] row : mat) {
      cur.clear();
      for (int a : pre) {
        for (int b : row) {
          cur.add(a + b);
        }
      }
      Collections.sort(cur);
      pre.clear();
      for (int i = 0; i < Math.min(k, cur.size()); ++i) {
        pre.add(cur.get(i));
      }
    }
    return pre.get(k - 1);
  }
}

class Solution {
public:
  int kthSmallest(vector<vector<int>>& mat, int k) {
    int pre[k];
    int cur[mat[0].size() * k];
    memset(pre, 0, sizeof pre);
    int size = 1;
    for (auto& row : mat) {
      int i = 0;
      for (int j = 0; j < size; ++j) {
        for (int& v : row) {
          cur[i++] = pre[j] + v;
        }
      }
      sort(cur, cur + i);
      size = min(i, k);
      for (int j = 0; j < size; ++j) {
        pre[j] = cur[j];
      }
    }
    return pre[k - 1];
  }
};

func kthSmallest(mat [][]int, k int) int {
  pre := []int{0}
  for _, row := range mat {
    cur := []int{}
    for _, a := range pre {
      for _, b := range row {
        cur = append(cur, a+b)
      }
    }
    sort.Ints(cur)
    pre = cur[:min(k, len(cur))]
  }
  return pre[k-1]
}

function kthSmallest(mat: number[][], k: number): number {
  let pre: number[] = [0];
  for (const cur of mat) {
    const next: number[] = [];
    for (const a of pre) {
      for (const b of cur) {
        next.push(a + b);
      }
    }
    pre = next.sort((a, b) => a - b).slice(0, k);
  }
  return pre[k - 1];
}