Vault 101

xjay2007's Blog

0%

LeetCode 378. 有序矩阵中第K小的元素

发表于 更新于 分类于 ， Valine：

题目描述

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

示例:

1
2
3
4
5
6
7
8
matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

返回 13。

说明:
你可以假设 k 的值永远是有效的, 1 ≤ k ≤ n2 。

题解

优先队列

算法

  • 定义 tuple<int,int,int> ，分别表示行、列和当前行列的值，即 {x,y,matrix[x][y]}。
  • 建立优先队列(最小堆)，将第一行所有元素插入队列。
  • 重复下列步骤 k-1 次：
    • 每次取出队列顶(最小值) top，找到行 x 与列 y。
    • 将 top 下一行相同列的元素 matrix[x+1][y] 放入队列。
  • 最后队列顶为第 k 小的元素。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
// C++
typedef tuple<int, int, int> Tuple;
int kthSmallest(vector<vector<int>>& matrix, int k) {
    int N = (int)matrix.size();
    auto cmp = [](Tuple& t1, Tuple& t2) {
        return get<2>(t1) > get<2>(t2);
    };
    priority_queue<Tuple, vector<Tuple>, decltype(cmp)> pq(cmp);
    for (int i = 0; i < N; ++i) {
        pq.push({0, i, matrix[0][i]});
    }
    for (int i = 0; i < k - 1; ++i) {
        auto t = pq.top();
        pq.pop();
        if (get<0>(t) == N - 1) {
            continue;
        }
        int x, y, val;
        tie(x, y, val) = t;
        pq.push({x + 1, y, matrix[x+1][y]});
    }
    return get<2>(pq.top());
}

复杂度分析

  • 时间复杂度：$O(N + k\log{N})$ 。建立优先队列 $O(N)$ ，搜索队列需要 $O(k \log{N})$ 。
  • 空间复杂度：$O(N)$ 。

优先队列 II

思路

遍历所有元素，如果元素 matrix[i][j] 小于队列顶，则插入队列，保持优先队列长度为 k ，最后队列顶为第 k 小。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// C++
int kthSmallest(vector<vector<int>>& matrix, int k) {
    int N = (int)matrix.size();
    priority_queue<int> pq;
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            if (pq.size() < k) {
                pq.push(matrix[i][j]);
            } else {
                if (pq.top() > matrix[i][j]) {
                    pq.push(matrix[i][j]);
                    pq.pop();
                }
            }
        }
    }
    return pq.top();
}

复杂度分析

  • 时间复杂度：$O(N^2\log(k))$ 。
  • 空间复杂度：$O(k)$ 。

二分查找

算法

思路是找出每次二分查找的搜索范围。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
// C++
int kthSmallest(vector<vector<int>>& matrix, int k) {
    int N = (int)matrix.size();
    int lo = matrix[0][0];
    int hi = matrix[N-1][N-1] + 1; // [lo, hi)
    while (lo < hi) {
        int mid = lo + (hi - lo)/2;
        int count = 0;
        int j = N - 1;
        for (int i = 0; i < N; ++i) {
            while (j >= 0 && matrix[i][j] > mid) {
                --j;
            }
            count += (j + 1);
        }
        if (count < k) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
}

复杂度分析

  • 时间复杂度：$O(N\log(hi - lo))$ 。
  • 空间复杂度：$O(1)$ 。

参考链接

相似文章