[NeetCode-LeetCode] Kth Largest Element in a Stream JAVA

반응형

728x90

[NeetCode-LeetCode] Kth Largest Element in a Stream - Easy

---

접근

• 힙 (우선순위 큐)

---

풀이

스트림에 있는 요소에 K번째로 큰 요소를 반환하는 문제이다. 처음 생성자를 호출하고 원소를 추가하는 add 연산만 구현하면 된다. add 연산 수행 시 요소를 추가하고 K 번째 요소가 어떤 것인지 반환하면 된다. 우선순위 큐의 문제라 적혀있던 만큼 우선순위 큐를 통해 구현하면 간단하게 풀 수 있는 문제이다.

class KthLargest {
    private PriorityQueue<Integer> minHeap;
    private int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        minHeap = new PriorityQueue<>();

        for (int num : nums) {
            add(num);
        }
    }

KthLargest 클래스를 구현해야 하고 생성자를 구현한 부분이다. 우선순위 큐를 선언하고 처음 넘겨받는 배열을 우선순위 큐에 넣어줘야 하는데 해당 과정이 add 연산을 통해 해결이 가능하기 때문에 add를 호출해 준다. 

    public int add(int val) {
        minHeap.offer(val);

        if (minHeap.size() > k) {
            minHeap.poll();
        }

        return minHeap.peek();
    }

add 연산으로 minHeap에 값을 추가하고 만약 현재 minHeap의 사이즈가 k보다 크다면 추가되었던 원소를 뽑는다. 이렇게 되면 minHeap은 항상 k와 동일한 사이즈로 유지되고 peek의 값이 k 번째 값이 되기 때문에 add 연산 당 삽입, 삭제가 최대 1번씩 이뤄지고 해결이 가능하다.

우선순위 큐

앞에 더 성능이 좋은 코드들이 있는데 저 코드들은 큐를 사용하지 않고 배열을 통해 큐의 동작 방식을 직접 구현한 코드들이었다. add 연산 수행 시 배열 내에 원소들을 직접 비교하여 우선순위를 맞춰주는 방식으로 구현이 되어 사실상 알고리즘 자체는 동일하기에 넘어갔다.

---

전체 코드

class KthLargest {
    private PriorityQueue<Integer> minHeap;
    private int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        minHeap = new PriorityQueue<>();

        for (int num : nums) {
            add(num);
        }
    }

    public int add(int val) {
        minHeap.offer(val);

        if (minHeap.size() > k) {
            minHeap.poll();
        }

        return minHeap.peek();
    }
}