有序矩阵中第K小元素[优先队列PriorityQueue]

前言

如果遇到求最短路径相关问题，而每走一步面临着n个选择方案，如果直接遍历O(n)不是想要的复杂度，则在N个元素中以O(logN)定位元素，非PriorityQueue莫属。

一、有序矩阵中第K个最小元素

二、PriorityQueue

package everyday.medium;

import java.util.Comparator;
import java.util.PriorityQueue;

// 有序矩阵中第K个最小元素。
public class KthSmallest {
    /*
    行有序/列有序，但不同行不同列之间没有关系。
    每个坐标可往下/右走，通过优先队列拿到最小坐标即可。
     */
    public int kthSmallest(int[][] matrix, int k) {
        int m = matrix.length;

        PriorityQueue<int[]> que = new PriorityQueue<>(Comparator.comparingInt(p -> p[2]));
        // bug2:修改数组标记后，原数据就不再了，需要将其记录。
        que.add(new int[]{0, 0, matrix[0][0]});
        // 标记
        matrix[0][0] = 1 << 31;

        while (!que.isEmpty()) {
            int[] pos = que.poll();
            int i = pos[0], j = pos[1];
            // 先判定是否寻找到了第K小元素。
            if (--k == 0) return pos[2];
            // bug1:元素入队列了之后需要标记，防止再入队列。
            if (i < m - 1 && matrix[i + 1][j] != 1 << 31) {
                que.add(new int[]{i + 1, j, matrix[i + 1][j]});
                // 标记
                matrix[i + 1][j] = 1 << 31;
            }
            if (j < m - 1 && matrix[i][j + 1] != 1 << 31) {
                que.add(new int[]{pos[0], pos[1] + 1, matrix[i][j + 1]});
                // 标记
                matrix[i][j + 1] = 1 << 31;
            }
        }
        return -1;
    }
}

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
33
34
35
36
37
38
39
40
41

总结

1）优先队列，以O(logN)的时间寻找元素，是求局部最小/最大的第时间复杂度方法。

参考文献

[1] LeetCode 有序矩阵中第K小元素