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

题目描述

输入一个 n x n 的矩阵 matrix，每一行从左到右、每一列从上到下都是非递减的。另给一个整数 k。

把矩阵所有元素排成一列从小到大排序，返回排在第 k 位的元素（k 从 1 开始数）。值相同的元素各占一个位置，不是求第 k 个不同的值。

要求使用的内存少于 O(n^2)。

示例 1

输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13

解释：排序后为 [1,5,9,10,11,12,13,13,15]，第 8 个是 13。

示例 2

输入：matrix = [[-5]], k = 1
输出：-5

数据范围

n == matrix.length == matrix[i].length
1 <= n <= 300
-10^9 <= matrix[i][j] <= 10^9
每行、每列都按非递减顺序排列
1 <= k <= n^2

进阶：能否只用 O(1) 的额外内存？能否做到 O(n) 时间？后者较难，可参考相关论文。

解题思路：值域二分 + 阶梯计数

对“值”二分，而不是对位置二分。答案在 [左上角, 右下角] 之间。对候选值 mid，数一数矩阵里有多少个数 <= mid：从右上角出发，当前格 > mid 就左移，否则这一行左边的 j + 1 个都 <= mid，计数后下移一行。计数 >= k 说明答案 <= mid，hi = mid；否则 lo = mid + 1。lo == hi 时就是答案，而且它一定是矩阵里的某个数。

为什么正确

第 k 小的元素始终在值区间 [lo, hi] 内；count(x) 随 x 单调不减。

复杂度分析

时间 O(n · log(max - min))，每次计数沿阶梯走 O(n) 步；空间 O(1)。

易错点

计数 >= k 时不能写 hi = mid - 1，mid 本身可能就是答案。值可以为负，C++/Java 里 (lo + hi) / 2 向 0 取整，lo=-3、hi=-2 时 mid=-2 会死循环，要用向下取整（右移）并防止 lo + hi 溢出。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        lo, hi = matrix[0][0], matrix[n - 1][n - 1]
        while lo < hi:
            mid = (lo + hi) // 2
            cnt, j = 0, n - 1
            for i in range(n):
                while j >= 0 and matrix[i][j] > mid:
                    j -= 1
                cnt += j + 1
            if cnt >= k:
                hi = mid
            else:
                lo = mid + 1
        return lo

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int lo = matrix[0][0], hi = matrix[n - 1][n - 1];
        while (lo < hi) {
            int mid = (int) (((long) lo + hi) >> 1);
            int cnt = 0, j = n - 1;
            for (int i = 0; i < n; i++) {
                while (j >= 0 && matrix[i][j] > mid) {
                    j--;
                }
                cnt += j + 1;
            }
            if (cnt >= k) {
                hi = mid;
            } else {
                lo = mid + 1;
            }
        }
        return lo;
    }
}

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        int lo = matrix[0][0], hi = matrix[n - 1][n - 1];
        while (lo < hi) {
            int mid = (int)(((long long)lo + hi) >> 1);
            int cnt = 0, j = n - 1;
            for (int i = 0; i < n; i++) {
                while (j >= 0 && matrix[i][j] > mid) {
                    j--;
                }
                cnt += j + 1;
            }
            if (cnt >= k) {
                hi = mid;
            } else {
                lo = mid + 1;
            }
        }
        return lo;
    }
};

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)
	lo, hi := matrix[0][0], matrix[n-1][n-1]
	for lo < hi {
		mid := (lo + hi) >> 1
		cnt, j := 0, n-1
		for i := 0; i < n; i++ {
			for j >= 0 && matrix[i][j] > mid {
				j--
			}
			cnt += j + 1
		}
		if cnt >= k {
			hi = mid
		} else {
			lo = mid + 1
		}
	}
	return lo
}

var kthSmallest = function(matrix, k) {
    const n = matrix.length;
    let lo = matrix[0][0], hi = matrix[n - 1][n - 1];
    while (lo < hi) {
        const mid = Math.floor((lo + hi) / 2);
        let cnt = 0, j = n - 1;
        for (let i = 0; i < n; i++) {
            while (j >= 0 && matrix[i][j] > mid) {
                j--;
            }
            cnt += j + 1;
        }
        if (cnt >= k) {
            hi = mid;
        } else {
            lo = mid + 1;
        }
    }
    return lo;
};

function kthSmallest(matrix: number[][], k: number): number {
    const n = matrix.length;
    let lo = matrix[0][0], hi = matrix[n - 1][n - 1];
    while (lo < hi) {
        const mid = Math.floor((lo + hi) / 2);
        let cnt = 0, j = n - 1;
        for (let i = 0; i < n; i++) {
            while (j >= 0 && matrix[i][j] > mid) {
                j--;
            }
            cnt += j + 1;
        }
        if (cnt >= k) {
            hi = mid;
        } else {
            lo = mid + 1;
        }
    }
    return lo;
}

载入编辑器，自己跑一遍 →

其他解法

直接排序不利用有序性；最小堆只利用行有序，代价与 k 相关；主解法同时利用行列有序，在值域上二分，空间 O(1)。

解法 2：最小堆多路归并时间 O(k log n)，空间 O(n)

每一行都是一个有序序列，问题变成合并 n 个有序序列后的第 k 个数。先把每行第一个元素和它的行列号放入最小堆。弹出 k - 1 次，每次把被弹出元素同一行的下一个元素入堆。此时堆顶就是答案。

只利用了行有序，思路通用，可直接用于合并 k 个有序链表一类题。k 接近 n² 时比主解法慢，并且需要 O(n) 的堆空间。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        heap = [(matrix[i][0], i, 0) for i in range(min(n, k))]
        heapq.heapify(heap)
        for _ in range(k - 1):
            v, i, j = heapq.heappop(heap)
            if j + 1 < n:
                heapq.heappush(heap, (matrix[i][j + 1], i, j + 1))
        return heap[0][0]

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
        for (int i = 0; i < Math.min(n, k); i++) heap.offer(new int[]{matrix[i][0], i, 0});
        for (int t = 0; t < k - 1; t++) {
            int[] top = heap.poll();
            int i = top[1], j = top[2];
            if (j + 1 < n) heap.offer(new int[]{matrix[i][j + 1], i, j + 1});
        }
        return heap.peek()[0];
    }
}

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<>> heap;
        for (int i = 0; i < min(n, k); i++) heap.emplace(matrix[i][0], i, 0);
        for (int t = 0; t < k - 1; t++) {
            auto [v, i, j] = heap.top();
            heap.pop();
            if (j + 1 < n) heap.emplace(matrix[i][j + 1], i, j + 1);
        }
        return get<0>(heap.top());
    }
};

type cellHeap [][3]int

func (h cellHeap) Len() int           { return len(h) }
func (h cellHeap) Less(a, b int) bool { return h[a][0] < h[b][0] }
func (h cellHeap) Swap(a, b int)      { h[a], h[b] = h[b], h[a] }
func (h *cellHeap) Push(x any)        { *h = append(*h, x.([3]int)) }
func (h *cellHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

func kthSmallest(matrix [][]int, k int) int {
	n := len(matrix)
	h := &cellHeap{}
	for i := 0; i < min(n, k); i++ {
		*h = append(*h, [3]int{matrix[i][0], i, 0})
	}
	heap.Init(h)
	for t := 0; t < k-1; t++ {
		top := heap.Pop(h).([3]int)
		i, j := top[1], top[2]
		if j+1 < n {
			heap.Push(h, [3]int{matrix[i][j+1], i, j + 1})
		}
	}
	return (*h)[0][0]
}

class MinHeap {
    constructor(less) { this.a = []; this.less = less; }
    get size() { return this.a.length; }
    top() { return this.a[0]; }
    push(x) {
        const a = this.a; a.push(x);
        for (let i = a.length - 1; i > 0;) {
            const p = (i - 1) >> 1;
            if (!this.less(a[i], a[p])) break;
            [a[i], a[p]] = [a[p], a[i]]; i = p;
        }
    }
    pop() {
        const a = this.a, top = a[0], last = a.pop();
        if (a.length) {
            a[0] = last;
            for (let i = 0; ;) {
                let m = i; const l = 2 * i + 1, r = l + 1;
                if (l < a.length && this.less(a[l], a[m])) m = l;
                if (r < a.length && this.less(a[r], a[m])) m = r;
                if (m === i) break;
                [a[i], a[m]] = [a[m], a[i]]; i = m;
            }
        }
        return top;
    }
}

var kthSmallest = function(matrix, k) {
    const n = matrix.length;
    const heap = new MinHeap((a, b) => a[0] < b[0]);
    for (let i = 0; i < Math.min(n, k); i++) heap.push([matrix[i][0], i, 0]);
    for (let t = 0; t < k - 1; t++) {
        const [v, i, j] = heap.pop();
        if (j + 1 < n) heap.push([matrix[i][j + 1], i, j + 1]);
    }
    return heap.top()[0];
};

class MinHeap<T> {
    private a: T[] = [];
    constructor(private less: (x: T, y: T) => boolean) {}
    get size(): number { return this.a.length; }
    top(): T { return this.a[0]; }
    push(x: T): void {
        const a = this.a; a.push(x);
        for (let i = a.length - 1; i > 0;) {
            const p = (i - 1) >> 1;
            if (!this.less(a[i], a[p])) break;
            [a[i], a[p]] = [a[p], a[i]]; i = p;
        }
    }
    pop(): T {
        const a = this.a, top = a[0], last = a.pop()!;
        if (a.length) {
            a[0] = last;
            for (let i = 0; ;) {
                let m = i; const l = 2 * i + 1, r = l + 1;
                if (l < a.length && this.less(a[l], a[m])) m = l;
                if (r < a.length && this.less(a[r], a[m])) m = r;
                if (m === i) break;
                [a[i], a[m]] = [a[m], a[i]]; i = m;
            }
        }
        return top;
    }
}

function kthSmallest(matrix: number[][], k: number): number {
    const n = matrix.length;
    const heap = new MinHeap<number[]>((a, b) => a[0] < b[0]);
    for (let i = 0; i < Math.min(n, k); i++) heap.push([matrix[i][0], i, 0]);
    for (let t = 0; t < k - 1; t++) {
        const [v, i, j] = heap.pop();
        if (j + 1 < n) heap.push([matrix[i][j + 1], i, j + 1]);
    }
    return heap.top()[0];
}

解法 3：直接排序时间 O(n² log n)，空间 O(n²)

把矩阵所有元素放进一个一维数组。排序后返回下标 k - 1 的元素。

实现最简单，完全没有利用矩阵有序的性质。时间和空间都最差，适合作为对拍基准。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        arr = [x for row in matrix for x in row]
        arr.sort()
        return arr[k - 1]

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int[] arr = new int[n * n];
        int p = 0;
        for (int[] row : matrix)
            for (int x : row) arr[p++] = x;
        Arrays.sort(arr);
        return arr[k - 1];
    }
}

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        vector<int> arr;
        for (auto& row : matrix)
            for (int x : row) arr.push_back(x);
        sort(arr.begin(), arr.end());
        return arr[k - 1];
    }
};

func kthSmallest(matrix [][]int, k int) int {
	arr := []int{}
	for _, row := range matrix {
		arr = append(arr, row...)
	}
	sort.Ints(arr)
	return arr[k-1]
}

var kthSmallest = function(matrix, k) {
    const arr = matrix.flat();
    arr.sort((a, b) => a - b);
    return arr[k - 1];
};

function kthSmallest(matrix: number[][], k: number): number {
    const arr = matrix.flat();
    arr.sort((a, b) => a - b);
    return arr[k - 1];
}

二分查找专题的其他题目

查看二分查找全部题目 →