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

难度: 🌕🌕

问题描述

---

解法 1 - PriorityQueue

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        // 思路：
        // 堆排
        // 将每列的头元素加入小根堆，每次弹出的都是最小的元素
        // 复杂度 O(klogM)
        int m = matrix.length;
        int n = matrix[0].length;
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> {
            return matrix[a[0]][a[1]] - matrix[b[0]][b[1]];
        });

        // 初始化
        for(int i = 0; i < m; i ++) {
            queue.offer(new int[] {i, 0});
        }

        int res = 0;
        for(int i = 0; i < k; i ++) {
            int[] cur = queue.poll();
            int x = cur[0];
            int y = cur[1];
            res = matrix[x][y];
            if(y + 1 < n) {
                queue.offer(new int[] {x, y + 1});
            }
        }
        return res;
    }
}

---

输出 1

---

解法 2 - 二分

class Solution {
    int m;
    int n;
    int[][] matrix;
    int k;
    public int kthSmallest(int[][] matrix, int k) {
        // 思路：
        // 这个二维数组有个规律，就是从右上角到左下角要找某个元素的路径的确定的
        // min = 左上角, max = 右下角元素
        // 通过二分法，每次确定一个 mid，按照右上 - 左下的路径可以将其分为两部分，一部分是小于等于它的，一部分是大于它的
        // 通过获取小于等于它的元素个数 count 与 k 的大小关系可以确定目标值在左半边还是右半边，从而继续二分直到找到目标值
        // 存在一个问题，就是如果找到的 mid 值在二维数组中不存在，但满足小于等于它的元素为 k 个，该情况如何解决？
        // 可以继续对左半边进行二分，直到区间只剩一个元素返回，那么这个元素在二维数组中肯定存在
        m = matrix.length;
        n = matrix[0].length;
        this.matrix = matrix;
        this.k = k;
        int min = matrix[0][0];
        int max = matrix[m - 1][n - 1];
        return mySol(min, max);
    }

    private int mySol(int left, int right) {
        // 递归终止条件
        if(left == right) {
            return left;
        }
        int mid = left + ((right - left) >> 1);
        int count = getCount(mid);
        // System.out.println(left + " " + right + "  " + mid + " : " + count);
        if(k == count) {
            return mySol(left, mid);
        } else if(k < count) {
            return mySol(left, mid); // 若包含多个 mid，那么结果可能为 mid，因此 mid 需要包含进去
        } else {
            return mySol(mid + 1, right);
        }
    }

    private int getCount(int target) {
        // 在二维数组中找到小于等于 target 的元素个数
        // 从右上角往左下角开始找
        // 复杂度 max = O(M + N)
        int res = 0;
        int i = 0;
        int j = n - 1;

        while(i < m && j >= 0) {
            // 先尽可能向下找，得到该列获得的元素个数
            while(i < m && target >= matrix[i][j]) {
                i ++;
            }
            res += i;
            if(i >= m) {
                // 说明该列所有元素均 < target，那么前面所有列肯定也 < target
                res += m * j;
                return res;
            } 
            // i < m 需要左移一列找前一列比 target 小于等于的元素个数
            j --;
        }
        return res;
    }
}

---

输出 2