[100天算法】-有序矩阵中第K小的元素（day 58）

题目描述

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
请注意，它是排序后的第 k 小元素，而不是第 k 个不同的元素。

 

示例：

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

返回 13。
 

提示：
你可以假设 k 的值永远是有效的，1 ≤ k ≤ n2 。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

方法1：归并排序

思路

矩阵的每一行都是一个升序的数组，用合并多个有序数组的思路，找出前 k 个最小的元素就行。

用 n 个指针来记录每一行当前最小元素。

复杂度分析

• 时间复杂度：$O(k*n)$，n 是矩阵宽高，其中 k 最坏的情况下是 n2，所以最坏的时间复杂度是 O(n3)。
• 空间复杂度：$O(n)$，指针数组的长度。

代码

JavaScript Code

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function (matrix, k) {
    const rows = matrix.length;
    const cols = matrix[0].length
    const indices = Array(rows).fill(0);

    let kth;
    while (k-- > 0) {
        let min = Infinity
        let minRow = -1
        for (let r = 0; r < rows; r++) {
            if (indices[r] >= cols) continue

            const num = matrix[r][indices[r]];
            if (num <= min) {
                min = num;
                minRow = r
            }
        }
        kth = min
        indices[minRow]++;
    }
    return kth;
};

方法2：归并排序+堆

思路

还是上一个方法的思路，只不过在寻找多行中的最小值时，不用循环查找，而是用堆将这个查找时间从 n 降到 logn。

复杂度分析

• 时间复杂度：$O(klogn)$，n 是矩阵宽高，其中 k 最坏的情况下是 $n^2$，所以最坏的时间复杂度是 $O(n^2logn)$。
• 空间复杂度：$O(n)$，堆的大小。

代码

JavaScript Code

/**
 * @param {number[][]} matrix
 * @param {number} k
 * @return {number}
 */
var kthSmallest = function (matrix, k) {
    const list = matrix.map((row, x) => [row[0], x, 0]);
    // 堆中的数据结构是 [num, x, y]
    const heap = new MinHeap(list, function comparator(inserted, compared) {
        return inserted[0] > compared[0];
    });
    
    while (k-- > 1) {
        const [, x, y] = heap.pop();
        if (y < matrix[0].length - 1) {
            heap.insert([matrix[x][y + 1], x, y + 1]);
        }
    }
    return heap.pop()[0];
};

// **************************************************

class Heap {
    constructor(list = [], comparator) {
        this.list = list;
        this.comparator = comparator;

        this.init();
    }

    init() {
        const size = this.size();
        for (let i = Math.floor(size / 2) - 1; i >= 0; i--) {
            this.heapify(this.list, size, i);
        }
    }

    insert(n) {
        this.list.push(n);
        const size = this.size();
        for (let i = Math.floor(size / 2) - 1; i >= 0; i--) {
            this.heapify(this.list, size, i);
        }
    }

    peek() {
        return this.list[0];
    }

    pop() {
        const last = this.list.pop();
        if (this.size() === 0) return last;
        const returnItem = this.list[0];
        this.list[0] = last;
        this.heapify(this.list, this.size(), 0);
        return returnItem;
    }

    size() {
        return this.list.length;
    }
}

class MinHeap extends Heap {
    constructor(list, comparator) {
        if (typeof comparator != 'function') {
            comparator = function comparator(inserted, compared) {
                return inserted > compared;
            };
        }
        super(list, comparator);
    }

    heapify(arr, size, i) {
        let smallest = i;
        const left = Math.floor(i * 2 + 1);
        const right = Math.floor(i * 2 + 2);
        if (left < size && this.comparator(arr[smallest], arr[left]))
            smallest = left;
        if (right < size && this.comparator(arr[smallest], arr[right]))
            smallest = right;

        if (smallest !== i) {
            [arr[smallest], arr[i]] = [arr[i], arr[smallest]];
            this.heapify(arr, size, smallest);
        }
    }
}