Sui Xin's Blog

[LeetCode]378. 有序矩阵中第K小的元素（Kth Smallest Element in a Sorted Matrix）

Posted on In Disqus:

题目描述

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

示例:

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

返回 13。

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

解题思路

  1. 解法一：对每一行第一个未加入的数字进行遍历，把最小的数字加入到结果中。时间复杂度较高。
  2. 解法二：使用最小优先队列，每次将刚加入结果的周围四个中未遍历过的元素加入到优先队列中，下次从优先队列中出队即为最小值。使用 Python 的 queue.PriorityQueue 来实现。

代码

Python 3.6

解法一

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        rows, cols = len(matrix), len(matrix[0])
        cur = [0] * cols
        max_num = matrix[-1][-1]
        count = 0
        while True:
            min_num = max_num
            min_idx = 0
            for i in range(rows):
                if cur[i] >= cols:
                    continue
                if matrix[i][cur[i]] < min_num:
                    min_num = matrix[i][cur[i]]
                    min_idx = i
            count += 1
            if count == k:
                return min_num
            cur[min_idx] += 1

执行用时：1852 ms
内存消耗：17.3 MB

解法二

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        rows, cols = len(matrix), len(matrix[0])
        if rows == 0 or cols == 0:
            return -1
        import queue
        next_step = [[0, 1], [1, 0], [0, -1], [-1, 0]]
        visited = [[0] * cols for _ in range(rows)]
        my_queue = queue.PriorityQueue()
        my_queue.put((matrix[0][0], 0, 0))
        visited[0][0] = 1
        step = 0
        while my_queue:
            cur_val, cur_x, cur_y = my_queue.get()
            step += 1
            if step == k:
                return cur_val
            for next_coor in next_step:
                next_x = cur_x + next_coor[0]
                next_y = cur_y + next_coor[1]
                if next_x < 0 or next_y < 0 or next_x >= rows or next_y >= cols or (visited[next_x][next_y] == 1):
                    continue
                my_queue.put((matrix[next_x][next_y], next_x, next_y))
                visited[next_x][next_y] = 1

执行用时：580 ms
内存消耗：17.7 MB

Related Posts