KK's blog

每天积累多一些

0%

LeetCode 378 Kth Smallest Element in a Sorted Matrix

Posted on Edited on Disqus:

LeetCode 378 Kth Smallest Element in a Sorted Matrix

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the k<sup>th</sup> smallest element in the matrix.

Note that it is the k<sup>th</sup> smallest element in the sorted order, not the k<sup>th</sup> distinct element.

You must find a solution with complexity better than O(n<sup>2</sup>).

Example 1:

<pre>Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8 Output: 13 Explanation: The elements in the matrix are [1,5,9,10,11,12,13,<u>13</u>,15], and the 8<sup>th</sup> smallest number is 13 </pre>

Example 2:

<pre>Input: matrix = [[-5]], k = 1 Output: -5 </pre>

Constraints:

  • n == matrix.length == matrix[i].length
  • 1 <= n <= 300
  • -10<sup>9</sup> <= matrix[i][j] <= 10<sup>9</sup>
  • All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
  • 1 <= k <= n<sup>2</sup>

Follow up: Could you solve the problem in O(n) time complexity?

题目大意：

按行按列有序矩阵中，找第k小的数。

Heap解题思路(推荐)：

见Heap知识点。 分别将(value, i, j)放入heap中，取出堆顶元素后，去(i, j)相邻右和下节点放入堆中。这个方法容易实现，所以推荐。

LeetCode 074 Search a 2D Matrix 每一行有序，下一行的首元素大于上一行的尾元素 + 找target
LeetCode 240 Search a 2D Matrix II 按行按列有序 + 找target
LeetCode 378 Kth Smallest Element in a Sorted Matrix 按行按列有序 + 找第k大
矩阵结构方面，第一道每一行都是独立，所以可以独立地按行按列做二分法
后两道，矩阵二维连续，所以解法都是类BFS，从某个点开始，然后比较它相邻的两个点。出发点不同，第二道在近似矩阵中点(右上角或左下角)，第三道在左上角出发。

注意事项：

  1. 将(value, i, j)放入heap中

Python代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def kthSmallest4(self, matrix: List[List[int]], k: int) -> int:
	OFFSETS = [(0, 1), (1, 0)]
	heap = [(matrix[0][0], 0, 0)]
	visited = set([(0, 0)])
	while heap:
		node = heapq.heappop(heap)
		k -= 1
		if k == 0:
			return node[0]
		for _dx, _dy in OFFSETS:
			x, y = node[1] + _dx, node[2] + _dy
			if x < 0 or x >= len(matrix) or y < 0 or y >= len(matrix[0]) or (x, y) in visited:
				continue
			heapq.heappush(heap, (matrix[x][y], x, y))
			visited.add((x, y))
	return -1

算法分析：

for循环是k个，每次循环理论上产生2个节点，所以总共是2^k个，复杂度为<code>O(klog(2<sup>k</sup>)</code>, 也就是<code>O(k<sup>2</sup>)</code>
由于矩阵最多n^2个元素，所以复杂度为<code>O(klog(n<sup>2</sup>)</code>
所以时间复杂度为O(klogn)，空间复杂度<code>O(n<sup>2</sup>)</code>。

数值二分法算法II解题思路：

第k的数运用数值二分法

解题步骤：

  1. 数值二分法
  2. 难点在于统计小于mid的个数。若遍历全矩阵比较慢，采用按行遍历，每行再用二分法找到小于mid的数的个数，再求和。

注意事项：

  1. 注意k--, k从1开始
  2. 每行统计小于mid个数用find smaller的模板。返回值要加1，因为下标转换成个数。

Python代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
	k -= 1
	if not matrix or not matrix[0]:
		return -1
	N, M = len(matrix), len(matrix[0])
	start, end, epsilon = matrix[0][0], matrix[N - 1][M - 1], 0.5
	while end - start > epsilon:
		mid = start + (end - start) / 2
		count = sum([self.get_count(matrix[i], mid) for i in range(N)])
		if k < count:
			end = mid
		elif k > count:
			start = mid
		else:
			start = mid
	return math.floor(end)

def get_count(self, nums: List[int], target: float) -> int:
	start, end = 0, len(nums) - 1
	while start + 1 < end:
		mid = start + (end - start) // 2
		if target > nums[mid]:
			start = mid
		elif target < nums[mid]:
			end = mid
		else:
			end = mid
	if nums[end] < target:
		return end + 1
	if nums[start] < target:
		return start + 1
	return 0

用bisect优化

Python代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
	k -= 1
	if not matrix or not matrix[0]:
		return -1
	N, M = len(matrix), len(matrix[0])
	start, end, epsilon = matrix[0][0], matrix[N - 1][M - 1], 0.5
	while end - start > epsilon:
		mid = start + (end - start) / 2
		count = sum([bisect.bisect(matrix[i], mid) + 1 for i in range(N)])
		if k < count:
			end = mid
		elif k > count:
			start = mid
		else:
			start = mid
	return math.floor(end)

算法分析：

while循环有log[(max - min)/epsilon]个，假设数字平均分布，复杂度是log(n), 每个循环按每行(n行)统计小于mid的个数，
每次统计调用get_count用了log(n)， 所以总时间复杂度为O(log(n) * nlogn)，空间复杂度O(1)。

Free mock interview