title
给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。
你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。
示例 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
提示：
m == mat.lengthn == mat.length[i]1 <= m, n <= 401 <= k <= min(200, n ^ m)1 <= mat[i][j] <= 5000mat[i]是一个非递减数组
solution
class Solution:
def kthSmallest(self, mat: List[List[int]], k: int) -> int:
if not mat or not mat[0]:
return 0
m, n = len(mat), len(mat[0])
hq = []
import heapq
import copy
visit = collections.defaultdict(bool)
pointers = [0 for _ in range(m)]
cur = sum([mat[i][0] for i in range(m)])
hq = [(cur, pointers)]
visit[tuple(pointers)] = True
res = []
for i in range(k):
top, pointers = heapq.heappop(hq)
res.append(top)
for j, point in enumerate(pointers):
if point < n - 1:
new_pointers = copy.copy(pointers)
new_pointers[j] += 1
if tuple(new_pointers) in visit:
continue
visit[tuple(new_pointers)] = True
heapq.heappush(hq, (top + mat[j][point + 1] - mat[j][point], new_pointers))
return res[-1]
二分法
dfs: 统计有多少个小于它的数量
class Solution:
def kthSmallest(self, mat: List[List[int]], k: int) -> int:
if not mat:
return 0
m, n = len(mat), len(mat[0])
cnt = 0
base = sum([l[0] for l in mat])
r = sum(l[-1] for l in mat)
def dfs(index, sumv, maxSum):
nonlocal cnt
if index >= m or sumv > maxSum or cnt > k:
return
dfs(index+1, sumv, maxSum)
for j in range(1, n):
tmp = sumv + mat[index][j] - mat[index][0]
if tmp > maxSum:
break
cnt += 1
dfs(index + 1, tmp, maxSum)
l = base
while l <= r:
mid = (l + r) // 2
cnt = 1
dfs(0, base, mid)
if cnt >= k:
r = mid - 1
else:
l = mid + 1
return l