1439.有序矩阵中的第-k-个最小数组和.md

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

/*
 * @lc app=leetcode.cn id=1439 lang=typescript
 *
 * [1439] 有序矩阵中的第 k 个最小数组和
 */

// @lc code=start
function kthSmallest1(mat: number[][], k: number): number {}
// @lc code=end

解法 1: 排序

function kthSmallest(mat: number[][], k: number): number {
  const m = mat.length,
    n = mat[0].length
  let nums = [0]
  for (let i = 0; i < m; i++) {
    const tmp: number[] = []
    for (let j = 0; j < n; j++) {
      for (let num of nums) {
        tmp.push(num + mat[i][j])
      }
    }
    nums = tmp.sort((a, b) => a - b).slice(0, k)
  }
  return nums[k - 1]
}

解法 2: 堆

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
}

Case

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