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

Categories:

Tags:

题目描述
思路分析
参考代码
相似题目

题目描述

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

思路分析

最小堆

二分查找

type Element struct {
	val int
	row int
	col int
}

type MinHeap []Element

func (h MinHeap) Len() int { return len(h) }

func (h MinHeap) Less(i, j int) bool { return h[i].val < h[j].val }

func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *MinHeap) Push(x interface{}) {
	*h = append(*h, x.(Element))
}

func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)
	h := &MinHeap{}
	heap.Init(h)

	// 将每行的第一个元素和其位置放入最小堆
	for i := 0; i < n; i++ {
		heap.Push(h, Element{val: matrix[i][0], row: i, col: 0})
	}

	// 弹出堆顶元素，直到弹出第 k 小的元素
	for i := 0; i < k-1; i++ {
		elem := heap.Pop(h).(Element)
		if elem.col+1 < n {
			// 如果当前行还有更多元素，加入下一个元素到堆中
			heap.Push(h, Element{val: matrix[elem.row][elem.col+1], row: elem.row, col: elem.col + 1})
		}
	}

	return heap.Pop(h).(Element).val
}

参考代码

func countLessEqual(matrix [][]int, mid int) int {
	count := 0
	n := len(matrix)
	row, col := n-1, 0
	for row >= 0 && col < n {
		if matrix[row][col] <= mid {
			count += row + 1
			col++
		} else {
			row--
		}
	}
	return count
}

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)
	left, right := matrix[0][0], matrix[n-1][n-1]

	for left < right {
		mid := left + (right-left)/2
		if countLessEqual(matrix, mid) < k {
			left = mid + 1
		} else {
			right = mid
		}
	}

	return left
}

➡️ 点击查看 Java 题解

write your code here

相似题目

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

719. 找出第 K 小的数对距离

668. 乘法表中第 k 小的数

本文作者: HGNULB

本文链接: https://hgnulb.github.io/blog/2025/32635930