[LeetCode] 378. Kth Smallest Element in a Sorted Matrix

ProblemPermalink

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

You must find a solution with a memory complexity better than O(n2).

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

Constraints:

• n == matrix.length == matrix[i].length
• 1 <= n <= 300
• -109 <= matrix[i][j] <= 109
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• 1 <= k <= n2

Follow up:

• Could you solve the problem with a constant memory (i.e., O(1) memory complexity)?
• Could you solve the problem in O(n) time complexity? The solution may be too advanced for an interview but you may find reading this paper fun.

SolutionPermalink

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1));
        for (int row = 0; row < matrix.length; row++) {
            for (int col = 0; col < matrix[0].length; col++) {
                maxHeap.offer(matrix[row][col]);
                if (maxHeap.size() > k) maxHeap.poll();
            }
        }
        return maxHeap.poll();
    }
}

ExplanationPermalink

• 우선순위 큐를 사용하면 크기별로 요소를 저장할 수 있고, 그 결과로 k번째로 작은 수를 찾을 수 있다. 우선순위 큐를 선언한다.

  PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1));
  

• 우선순위 큐에 배열 matrix의 요소들을 저장하면서 k번째로 저장된 요소 이후의 요소들은 모두 빼버림으로써 k번째로 들어온 요소가 가장 높은 우선순위를 가지도록 한다.

  for (int row = 0; row < matrix.length; row++) {
    for (int col = 0; col < matrix[0].length; col++) {
        maxHeap.offer(matrix[row][col]);
        if (maxHeap.size() > k) maxHeap.poll();
    }
  }
  

• 우선순위 큐에 저장된 가장 높은 우선순위의 요소를 반환한다.

  return maxHeap.poll();