[Grind75-LeetCode] Find Median from Data Stream JAVA

반응형

728x90

[Grind75-LeetCode] Find Median from Data Stream - Hard 

---

접근

• heap (우선순위 큐)

---

풀이

중앙값을 찾는 클래스를 구현하는 문제이다. 중앙값이란 주어진 원소들을 순서대로 나열했을 때 중앙에 위치한 값을 말한다. 홀수인 경우 중앙값은 하나이지만 짝수인 경우 한 개의 중앙값이 존재하는 것이 아닌 중앙에 위치한 두 개의 값을 더하여 2로 나눈 값이 중앙값이 된다. 1,2,3,4 가 있다면 중앙의 2,3을 더한 값인 5에서 2로 나눈 2.5가 중앙값이 된다. 이러한 중앙값을 찾을 수 있는 클래스를 구현해야 하는데 정수를 추가하는 addNum과 중앙값을 반환하는 findMedian 두 메서드 구현이 목표이다.

 

이 문제는 간단하게 생각하면 배열을 두고 길이를 통해 중앙값 인덱스의 위치를 계산해 반환하면 되지만 문제는 요소가 정렬된 상태에서 중앙값을 찾아야 한다는 소리다. 그렇다면 매번 중앙값을 찾기 전 배열을 정렬해주어야 하고 정렬엔 많은 시간이 소요되어 수많은 연산이 일어나는 테스트케이스에선 시간 초과가 나게 된다. 

class MedianFinder {
    private PriorityQueue<Integer> maxHeap;
    private PriorityQueue<Integer> minHeap;

    public MedianFinder() {
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        minHeap = new PriorityQueue<>();
    }

힙을 통해 구현된 우선순위 큐를 사용해 요소들을 관리해 주면 된다. 두 개의 힙을 사용하는데 하나는 내림차순으로 정렬된 maxHeap, 하나는 오름차순으로 정렬된 minHeap을 사용해 준다. 두 개의 힙을 사용하며 입력으로 받은 요소들을 절반씩 나눠 담아준다. maxHeap의 경우 절반을 기준으로 작은 값을 넣고 minHeap의 경우 절반을 기준으로 큰 값을 넣는 것이다. 

크기가 9인 배열이 존재한다면 [0:4]까지 5개의 요소는 maxHeap에 [5:8]까지 4개의 요소는 minHeap에 넣게 된다. maxHeap이 minHeap보다 무조건 크거나 같게 관리하며 짝수와 홀수의 경우를 구분해 주고 짝수인 경우 각각 힙의 peek 값을 더하여 나눠주게 되면 중앙값이 된다. maxHeap은 절반 중 작은 값이 내림차순으로 정렬되어 [0:4] 중 제일 큰 값인 4가 peek로 오게 되고 minHeap은 절반 중 큰 값이 오름차순으로 정렬되어 [5:8] 중 제일 작은 값인 5가 peek로 오게 된다.

    public void addNum(int num) {
        maxHeap.offer(num);

        if (!maxHeap.isEmpty() && !minHeap.isEmpty() && maxHeap.peek() > minHeap.peek()) {
            minHeap.offer(maxHeap.poll());
        }

        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.offer(maxHeap.poll());
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }

숫자를 추가하는 메서드이다. 우선 작은 값을 담는 큐로 먼저 추가하게 된다. 홀수의 경우 maxHeap에서 중앙값을 반환하게 구현하였기 때문에 maxHeap에 먼저 추가하게 되는 것이다. 그리고 maxHeap과 minHeap에 요소가 있고 maxHeap의 peek가 minHeap의 peek보다 크다면 요소를 minHeap으로 옮기게 된다. 이 과정은 두 힙 간 정렬을 위해 올바른 힙으로 요소를 옮기는 과정이다. 정렬이 끝났다면 힙의 사이즈를 관리해주어야 하는데 maxHeap을 minHeap보다 크거나 같게 유지하게 위해 옮겨 주는 과정이다.

    public double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        } else {
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        }
    }

중앙값을 반환하는 메서드이다. maxHeap이 minHeap보다 큰 경우 즉 홀수인 경우 maxHeap의 peek가 중앙값이 된다. 그리고 크기가 같은 짝수인 경우 두 힙의 peek 값을 더해 2로 나눈 뒤 반환해 준다.

 

힙

 

---

전체 코드

class MedianFinder {
    private PriorityQueue<Integer> maxHeap;
    private PriorityQueue<Integer> minHeap;

    public MedianFinder() {
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        minHeap = new PriorityQueue<>();
    }

    public void addNum(int num) {
        maxHeap.offer(num);

        if (!maxHeap.isEmpty() && !minHeap.isEmpty() && maxHeap.peek() > minHeap.peek()) {
            minHeap.offer(maxHeap.poll());
        }

        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.offer(maxHeap.poll());
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }

    public double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        } else {
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        }
    }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */