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

122 阅读1分钟

携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第46天,有序矩阵中第K小元素[优先队列PriorityQueue] - 掘金 (juejin.cn)

前言

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

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

image.png

二、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)优先队列,以O(logN)的时间寻找元素,是求局部最小/最大的第时间复杂度方法。

参考文献

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