LeetCode/Heap

[Medium] 378. Kth Smallest Element in a Sorted Matrix (아마존)

Developer07 2024. 4. 3. 14:02

#문제 설명

378번 문제는 정수 k와 n x n 정사각형 matrix가 주어지고, matrix의 각 칸에는 정수들이 담겨있다. 이 문제에서는 matrix안에서 정수 k번째로 작은 정수들을 전부 return해주면 된다. 예를 들어 [[1,5,9], [10,11,13] , [12,13,15]]  라는 2d 매트릭스가 주어지고 정수 k = 4가 주어지면, matrix안의 4번째로 작은 10을 return해주면 된다.

#문제 풀이

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(Collections.reverseOrder());

        for(int i = 0; i < matrix.length;i++){
            for(int j = 0; j < matrix[0].length;j++){
                maxHeap.add(matrix[i][j]);
                if(maxHeap.size() > k){
                    maxHeap.poll();
                }

            }
        }
        return maxHeap.peek();
    }
}

이 문제는 maxHeap을 사용해 풀 수 있다. matrix의 각 칸을 한번씩 loop해주면서 정수들을 maxHeap에 담고, heap의 크기가 k보다 적을시, 현재 힙에서 가장 큰 수는 필요 없어지기 때문에 가장 큰 수, 즉 root의 값을 poll 해주면 된다. 그리고 마지막에 maxHeap을 peek, 즉 root값을 return시켜주면 된다. 이 경우에 root는 matrix에서 k번째로 작은 수이고 N이 matrix의 크기라 할때 N - k번째로 큰 수를 return하게 되는 것이다.