function kthSmallest1(mat: number[][], k: number): number {
const m = mat.length,
n = mat[0].length
const queue = new Heap<[number, number[]]>((a, b) => a[0] - b[0])
queue.push([mat.reduce((a, b) => a + b[0], 0), new Array(m).fill(0)])
const vis = new Set<string>()
while (k-- > 1) {
const [sum, ids] = queue.pop()!
for (let i = 0; i < m; i++) {
const j = ids[i]
if (j < n - 1) {
ids[i] = j + 1
const key = ids.join(',')
if (!vis.has(key)) {
queue.push([sum + mat[i][j + 1] - mat[i][j], [...ids]])
vis.add(key)
}
ids[i] = j
}
}
}
const [sum] = queue.pop()
return sum
}
class Heap<T = number> {
private _heap: T[] = []
constructor(...args: T extends number ? [((a: number, b: number) => number)?, T[]?] : [(a: T, b: T) => number, T[]?])
constructor(private _comparator: (n1: T, n2: T) => number = ((a: number, b: number) => a - b) as any, arr?: T[]) {
if (arr) {
this._heap = [...arr]
for (let i = arr.length >> 1; i; i--) this.down(i)
}
}
private swap(i: number, j: number) {
;[this._heap[i], this._heap[j]] = [this._heap[j], this._heap[i]]
}
private _has(i: number) {
return Object.prototype.hasOwnProperty.call(this._heap, i)
}
private comparator(i: number, j: number) {
return this._comparator(this._heap[i], this._heap[j]) < 0
}
private parent(i: number) {
return (i - 1) >> 1
}
private children(i: number) {
return [2 * i + 1, 2 * i + 2]
}
private up(i: number) {
let parent = this.parent(i)
while (i > 0 && this.comparator(i, parent)) {
this.swap(i, parent)
;[i, parent] = [parent, (parent - 1) >> 1]
}
}
private down(i: number) {
let [left, right] = this.children(i)
while (left < this._heap.length) {
if (!this._has(right)) right = left
if (this.comparator(right, left)) [left, right] = [right, left]
if (this.comparator(i, left)) break
this.swap(i, left)
i = left
;[left, right] = [2 * i + 1, 2 * i + 2]
}
}
push(node: T) {
this._heap.push(node)
this.up(this._heap.length - 1)
}
pop() {
if (this.size() === 0) return null
let res = this._heap[0]
this.swap(0, this._heap.length - 1)
this._heap.pop()
this.down(0)
return res
}
remove(i: number) {
if (i >= this.size()) return
let res = this._heap[i]
this.swap(i, this._heap.length - 1)
this._heap.pop()
this.down(i)
this.up(i)
return res
}
modify(i: number, val: T) {
if (i >= this.size()) return
let res = this._heap[i]
this._heap[i] = val
this.down(i)
this.up(i)
return res
}
top() {
if (this.size() === 0) return null
return this._heap[0]
}
size() {
return this._heap.length
}
}
function kthSmallest(mat: number[][], k: number): number {
const m = mat.length,
n = mat[0].length
let min = 0,
max = 0
for (let i = 0; i < m; i++) {
min += mat[i][0]
max += mat[i][n - 1]
}
const check = (t: number) => {
let res = 1
const dfs = (i: number, sum: number) => {
if (sum > t) return
if (i === m) return
if (res > k) return
dfs(i + 1, sum)
for (let j = 1; j < n; j++) {
const a = sum + mat[i][j] - mat[i][0]
if (a > t) break
res++
dfs(i + 1, a)
}
}
dfs(0, min)
return res >= k
}
let l = min,
r = max
while (l < r) {
let m = (l + r) >> 1
if (check(m)) {
r = m
} else {
l = m + 1
}
}
return l
}
test.each([
{
input: {
mat: [
[13, 88, 148, 211, 300, 330, 399],
[71, 123, 163, 229, 289, 346, 357],
[45, 47, 49, 75, 202, 241, 283],
[45, 48, 146, 231, 243, 372, 400],
[40, 192, 271, 279, 285, 308, 368],
[128, 137, 173, 221, 344, 361, 368],
[67, 107, 119, 281, 372, 384, 396],
[2, 6, 78, 102, 230, 265, 355],
[2, 69, 97, 134, 157, 331, 392],
[77, 147, 175, 213, 248, 336, 355],
[7, 58, 202, 275, 283, 339, 366],
[74, 101, 158, 162, 330, 363, 371],
[78, 138, 177, 177, 197, 204, 379],
],
k: 50,
},
output: 691,
},
{
input: {
mat: [
[1, 3, 11],
[2, 4, 6],
],
k: 5,
},
output: 7,
},
{
input: {
mat: [
[1, 3, 11],
[2, 4, 6],
],
k: 9,
},
output: 17,
},
{
input: {
mat: [
[1, 10, 10],
[1, 4, 5],
[2, 3, 6],
],
k: 7,
},
output: 9,
},
])('input: mat = $input.mat, k = $input.k', ({ input: { mat, k }, output }) => {
expect(kthSmallest(mat, k)).toEqual(output)
})