前端就该用 JS 刷算法32

323 阅读2分钟

每日一题 -- 堆

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

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

分析

  1. 数组都是有序的,用 len 个指针指向每一个数组的下标
  2. 最小堆维护一个最小的总值和它对应的指针数组,每一次出堆的时候,根据它的指针数组,分别对每一个指针向前1位,得到新的指针和数组,再放回堆中

注意

  1. 由于指针数组的动态变更的,但是[1,0] -> [1,1] 和 [0,1] -> [1,1] 得到的结果是一样的,所以要用 map 保证每次入堆的指针是唯一的
  2. 由于 map 的 key 不能是数组,所以 map.set 的时候要将数组转成字符串
  3. 每当出堆的时候的指针数组,需要分别+1,所以需要对数组做一次拷贝
// [1439. 有序矩阵中的第 k 个最小数组和](https://leetcode-cn.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/)

/**
 * @分析
 * 1. 初始化指针数组,和对应的数组和
 */
var kthSmallest = function (mat, k) {
    const len = mat.length // 矩阵的长度
    const arrLen = mat[0].length // 一维数组的长度
    let points = new Array(len).fill(0)
    const minHeap = new Heap()
    minHeap.heappush(points.reduce((pre, val, index) => pre + mat[index][val], 0), points)
    const map = new Map()
    map.set(arrToString(points),1)
    while (--k) {
        // 弹出 k-1 次就是第 k-1 小,最后一次不需要再处理,直接返回
        const [sum, points] = minHeap.heappop()
        for (let i = 0; i < len; i++) {
            if (points[i] < arrLen-1) {
                const temp =[...points]
                // 最长不能超出一维数组的长度
                temp[i] += 1
                if(!map.has(arrToString(temp))){
                    minHeap.heappush(temp.reduce((pre, val, index) => pre + mat[index][val], 0), temp)
                    map.set(arrToString(temp),1)
                }
            }
        }
    }
    return minHeap.data[1][0]
};

const arrToString = arr => {
    return JSON.stringify(arr)
}


/**
 * 实现小顶堆
 */

const Heap = function () {
    this.data = []
    // 第一个值是堆的实际长度,也就是堆末的下标
    this.data.push(0)
}


//  追加一个值
Heap.prototype.heappush = function (val, points = []) {
    // 往尾部加一个值,然后 up 上去
    this.data.push([val,points])
    this.data[0] += 1
    this.up(this.data[0])
}

//  弹出堆顶
Heap.prototype.heappop = function () {
    // 先将堆顶和堆最后一个值交换,删除,然后down 下来
    if(this.data[0] !== 1){
        this.swap(1, this.data[0])
    }
    const res = this.data.pop()
    this.data[0] -= 1
    this.down() // 默认就是从 1 开始
    return res // 将 pop 出来的值保存一下
}

Heap.prototype.swap = function (a, b) {
    let  temp = this.data[a] 
    this.data[a]  = this.data[b] 
    this.data[b] = temp
}

// 这边第一个值就当是当前堆的长度好了 -- 默认是从根节点开始
Heap.prototype.down = function (index = 1) {
    // 已经到最后一个可以找到左右节点的第二层节点了,这样 left 和 right 就不需要再判断了
    if (index * 2 > this.data[0]) return
    const left = 2 * index
    const right = 2 * index + 1
    let target = index
    if (left <= this.data[0] && this.data[left][0] < this.data[target][0]) {
        target = left
    }
    if (right <= this.data[0] && this.data[right][0] < this.data[target][0]) {
        target = right
    }
    if (target !== index) {
        this.swap(target, index)
        // 换的是值,还得继续往下面走
        this.down(target)
    }
    // 如果没变,证明走到这里已经整理完了,下面的子树已经是ok的了
}

Heap.prototype.up = function (index) {
    if (index < 2) return
    const fatherIndex = Math.floor(index / 2)
    // 只需要和父节点比较,兄弟节点是没有比较的价值的
    if (this.data[index][0] < this.data[fatherIndex][0]) {
        this.swap(index, fatherIndex)
        this.up(fatherIndex)
    }
}