LintCode 401. Kth Smallest Number in Sorted Matrix (堆经典好题!!!)

原创已于 2022-11-30 15:06:00 修改 · 518 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#LintCode

于 2019-01-19 16:34:30 首次发布

algorithm-design 专栏收录该内容

817 篇文章

订阅专栏

本文探讨了在排序矩阵中寻找第K小元素的有效算法，包括使用最大堆、最小堆和二分法等多种方法，详细解析了不同算法的时间复杂度及适用场景。

Kth Smallest Number in Sorted Matrix
Find the kth smallest number in a row and column sorted matrix.

Example
Given k = 4 and a matrix:

[
[1 ,5 ,7],
[3 ,7 ,8],
[4 ,8 ,9],
]
return 5

Challenge
Solve it in O(k log n) time where n is the bigger one between row size and column size.

解法1：用最大堆。即不管三七二十一，将矩阵元素塞进最大堆中。然后只要堆的size > k，就pop()，此时pop的是当前堆的最大元素。显然，前面mn-k大的元素都被pop掉了(一共pop了mn-k次)，剩下的元素一共k个，即k个最小的元素，其中top即为第k小的元素。
这个方法简单，但是没有用上matrix行列都有序这个条件。时间复杂度O(MNlogMN)。
代码如下：

class Solution {
public:
    /**
     * @param matrix: a matrix of integers
     * @param k: An integer
     * @return: the kth smallest number in the matrix
     */
    int kthSmallest(vector<vector<int>> &matrix, int k) {
        
        priority_queue<int> pq;
        int nRow = matrix.size();
        int nCol = matrix[0].size();
        
        for (int i = 0; i < nRow; ++i) {
            for (int j = 0; j < nCol; ++j) {
                pq.push(matrix[i][j]);
                if (pq.size() > k) pq.pop();
            }
        }
        
        return pq.top();
    }
};

解法2：用最大堆。
这个方法可以用于行列皆没有排序的matrix。该方法不是最优，因为没有用上每列都有序这个条件。

注意：如果没有行列皆排序这个条件，那么，求第K大的数我们应该用最小堆，求第K小的数我们应该用最大堆。
该方法的时间复杂度为O(MNlogM)。

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        if (n == 0) return 0;
        int m = matrix[0].size();
        priority_queue<int> maxHeap;
        vector<int> pos(n, 0);
        int count = 0, total_count = m * n;
        while (count < total_count) {
            for (int i = 0; i < n; ++i) {
                int candidate = matrix[i][pos[i]];
                if (count < k) {
                    maxHeap.push(candidate);
                    pos[i]++;
                }
                else if (candidate < maxHeap.top()) {
                    maxHeap.push(candidate);
                    maxHeap.pop();
                    pos[i]++;
                }
                count++;
            }
        }
        return maxHeap.top();
    }
};

解法3：用最小堆。
思路跟解法1有点像，但是反过来:
解法1是踢掉前MN-K个最大的值，用maxheap。
解法3是踢掉前K-1个最小的值，用minHeap。

我们建最小堆，先放matrix[0][0]，即最小的那个元素，然后将其pop()，再把它的正下方和右边的那个元素放进堆，因为这两个肯定是仅比top大的那两个最小的元素。如此反复，一共pop掉k-1次，并且每次pop掉top后都把它正下方和右方的元素加进来。
显然，k-1个最小的元素都被pop掉了，现在堆的top是第k小的元素。
时间复杂度应该是klogk。

struct Elem {
    int r;
    int c;
    int v;
    Elem(int row = 0, int col = 0, int val = 0) : r(row), c(col), v(val) {}

};

struct cmp {
    bool operator() (const Elem & a, const Elem &b) {
        return (a.v > b.v);
    }
};

class Solution {
public:
    /**
     * @param matrix: a matrix of integers
     * @param k: An integer
     * @return: the kth smallest number in the matrix
     */
    int kthSmallest(vector<vector<int>> &matrix, int k) {
        
        priority_queue<Elem, vector<Elem>, cmp> pq;  //min-heap
        int nRow = matrix.size();
        int nCol = matrix[0].size();
        
        pq.push(Elem(0, 0, matrix[0][0]));
        vector<vector<int>> visited(nRow, vector<int>(nCol, 0));
        visited[0][0] = 1;
        
        for (int i = 0; i < k - 1; ++i) {
            Elem elem = pq.top();  //totally popped k - 1 elements
            pq.pop();
            
            int newR = elem.r + 1;
            if (newR < nRow && !visited[newR][elem.c]) {
                pq.push(Elem(newR, elem.c, matrix[newR][elem.c]));
                visited[newR][elem.c] = 1;
            }
            
            int newC = elem.c + 1;
            if (newC < nCol && !visited[elem.r][newC]) {
                pq.push(Elem(elem.r, newC, matrix[elem.r][newC]));
                visited[elem.r][newC] = 1;
            }
        }
        
        return pq.top().v;
    }
};

二刷:
注意：重要!
visited[][]一定要设计成nRow和nCol的二维矩阵，不要设计成set visited。
否则

如果operator<()里面是return val > a.val; 那么同值的不同Node进set两次，会被判断成相同的两个节点
以输入
[[1,5,9],[10,11,13],[12,13,15]]
8
为例，里面的两个13会判断成为相同的两个节点。
如果operator<()里面是return val >= a.val; 那么同一个Node进set两次，会被判断成不一样的两个节点。

代码如下：

struct Node {
    int row;
    int col;
    int val;
    Node(int x, int y, int v) : row(x), col(y), val(v) {}
    bool operator < (const Node & a) const {
        return val > a.val;
    }
};

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int nRow = matrix.size();
        int nCol = matrix[0].size();
        priority_queue<Node> minHeap;
        vector<vector<int>> visited(nRow, vector<int>(nCol, 0));
        Node node = Node(0, 0, matrix[0][0]);
        minHeap.push(node);
        visited[0][0] = 1;
        int index = 1;
        while (!minHeap.empty()) {
            Node topNode = minHeap.top();
            minHeap.pop();
            if (index == k) return topNode.val;
            if (topNode.row + 1 < nRow) {
                Node newNode = Node(topNode.row + 1, topNode.col, matrix[topNode.row + 1][topNode.col]);
                if (visited[topNode.row + 1][topNode.col] == 0) {
                    minHeap.push(newNode);
                    visited[topNode.row + 1][topNode.col] = 1;
                }
            }
            if (topNode.col + 1 < nCol) {
                Node newNode = Node(topNode.row, topNode.col + 1, matrix[topNode.row][topNode.col + 1]);
                if (visited[topNode.row][topNode.col + 1] == 0) {
                    minHeap.push(newNode);
                    visited[topNode.row][topNode.col + 1] = 1;
                }
            }
            index++;
        }
        return -1;
    }
};

解法4：还是最小堆，可以直接
用LintCode 1874: Kth Smallest Element in a Specific Array中的做法，
https://blog.youkuaiyun.com/roufoo/article/details/105308293
一开始把每个有效行的头元素放进堆，然后每次pop后取pop节点的右节点。该算法时间复杂度是O(klogM) M为有效行数。有效是指该行不为空。这个算法没用到列也有序这个条件。
那么我们可以看出，如果M<k时，似乎该算法更好?

代码如下：

struct Node {
    int x;
    int y;
    int v;
    Node(int _x = 0, int _y = 0, int _v = 0) : x(_x), y(_y), v(_v) {}
    bool operator < (const Node & a) const {
        return v > a.v;
    }
};
 
class Solution {
public:
    /**
     * @param arr: an array of integers
     * @param k: an integer
     * @return: the Kth smallest element in a specific array
     */
    int kthSmallest(vector<vector<int>> &arr, int k) {
        int rowNum = arr.size();
        priority_queue<Node> pq;
        
        for (int i = 0; i < rowNum; ++i) {
            if (arr[i].size() > 0) {
                pq.push(Node(i, 0, arr[i][0]));
            }
        }
        int count = 0;
        while(!pq.empty()) {
            Node top = pq.top();
            pq.pop();
            count++;
            if (count == k) return top.v;
            if (top.y < arr[top.x].size() - 1) {
                pq.push(Node(top.x, top.y + 1, arr[top.x][top.y + 1]));
            }
        }
        return -1;
    }
};