378. Kth Smallest Element in a Sorted Matrix

QuickSort

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int targetIndex = k - 1;

        return quickSort(matrix, 0, n * n - 1, targetIndex);
    }

    private int quickSort(int[][] matrix, int left, int right, int targetIndex) {
        int n = matrix.length;
        int pivotIndex = partition(matrix, left, right);
        if (pivotIndex < targetIndex) {
            return quickSort(matrix, pivotIndex + 1, right, targetIndex);
        } else if (pivotIndex > targetIndex) {
            return quickSort(matrix, left, pivotIndex - 1, targetIndex);
        } else {
            return matrix[pivotIndex / n][pivotIndex % n];
        }
    }

    private int partition(int[][] matrix, int left, int right) {
        int n = matrix.length;
        int pivotValue = matrix[left / n][left % n];
        int i = left, j = right;
        while (i < j) {
            while (i < j && matrix[j / n][j % n] >= pivotValue) {
                j--;
            }
            while (i < j && matrix[i / n][i % n] <= pivotValue) {
                i++;
            }

            swap(matrix, i, j);
        }

        swap(matrix, i, left); // Don't forget this swap

        // now i equals j
        return i;
    }

    private void swap(int[][] matrix, int i, int j) {
        int n = matrix.length;

        int tmp = matrix[j / n][j % n];
        matrix[j / n][j % n] = matrix[i / n][i % n];
        matrix[i / n][i % n] = tmp;
    }
}

MergeSort

class Solution {
    private static class Cell implements Comparable {
        private final int value;
        private final int i;
        private final int j;

        public Cell(int value, int i, int j) {
            this.value = value;
            this.i = i;
            this.j = j;
        }

        @Override
        public int compareTo(Object o) {
            return this.value - ((Cell) o).value;
        }
    }

    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Cell> minHeap = new PriorityQueue<>();

        for (int i = 0; i < matrix.length; i++) {
            minHeap.offer(new Cell(matrix[i][0], i, 0));
        }

        // 取出 k - 1 个最小的元素，保证堆顶元素为第 k 小的元素
        for (int i = 1; i <= k - 1; i++) {
            Cell cell = minHeap.poll();
            if (cell.j < matrix.length - 1) {
                minHeap.offer(new Cell(matrix[cell.i][cell.j + 1], cell.i, cell.j + 1));
            }
        }

        return minHeap.poll().value;
    }
}

Binary Search

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int left = matrix[0][0], right = matrix[n - 1][n - 1];
        while (left < right) {
            int mid = (left + right) >> 1; // 注意此处 left 和 right 可能为负数，不能使用无符号右移

            // 计算小于等于 mid 的元素个数
            int lessThanOrEqualToNumCount = getLessThanOrEqualToNumCount(matrix, mid);
            if (lessThanOrEqualToNumCount < k) {
                left = mid + 1;
            } else if (lessThanOrEqualToNumCount > k) {
                right = mid; // 小于等于 mid 的元素数量大于 k, 但是因为可能存在重复元素，所以我们不能使用 right = mid - 1, 而是使用 right = mid
            } else {
                right = mid; // 不能使用 right = mid - 1 的原因同上
            }
        }

        // now left and right are equal
        return left;
    }

    private int getLessThanOrEqualToNumCount(int[][] matrix, int mid) {
        int n = matrix.length;
        int i = n - 1, j = 0;
        int count = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= mid) {
                count += i + 1;
                j++;
            } else {
                i--;
            } 
        }

        return count;
    }
}

References

378. Kth Smallest Element in a Sorted Matrix
此题二分查找过程的动态演示及分析