(Heap, Medium) Find Median from Data Stream

https://leetcode.com/problems/find-median-from-data-stream/description/

문제 구성만 보면 PriorityQueue 사용해서 어떻게든 풀 수 있는 것 같은데..
비효율적이지만 일단 정확성만 생각하고 아래같이 풀어봤다.

class MedianFinder {

    PriorityQueue<Integer> pq;

    public MedianFinder() {
        pq = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        pq.offer(num);
    }
    
    public double findMedian() {
        int size = pq.size();

        int mid = size / 2;
        if (size % 2 == 0) {
            mid--;
        }

        List<Integer> temp = new ArrayList<>();

        int index = 0;
        while (!pq.isEmpty() && index < mid) {
            temp.add(pq.poll());
            index++;
        }

        double result;
        if (size % 2 == 0) {
            int m1 = pq.poll();
            int m2 = pq.poll();
            
            result = (m1 + m2) / 2.0;
            temp.add(m1);
            temp.add(m2);
        } else {
            int m = pq.poll();
            result = (double) m;
            temp.add(m); 
        }

        pq.addAll(temp);

        return result;
    }
}

당연히 테스트케이스가 커지면 문제가 생김..
원본 PQ에 계속해서 연산 시도하고 있기 때문

그럼 좋은 아이디어가 뭐냐?
힙 두개를 이용해서 left, right를 구성하고,
각각은 Max Heap, Min Heap으로 구성한 뒤 peek() 하면 O(1)으로 해결할 수 있다.
(물론 넣는 작업 자체는 O(log N)임)

이렇게 하기 위해서 addNum 작업 시 두 PQ에 균등하게 넣어주는게 중요하다.
왼쪽 (MaxHeap)을 우선으로 하고, 현재 num값이 left로 갈지 right로 갈지 peek()을 통해 판단한다.

이후 더 중요한건 리밸런싱인데, left 사이즈와 right 사이즈가 2이상 차이나면 서로에게서 poll한다음 옮겨준다.

모든 작업 시 고려해야 할 점은 left가 우선권이 있기 때문에 사이즈가 같거나, 혹은 right보다 1 더 클 수 있다는 점이다.

class MedianFinder {

    private PriorityQueue<Integer> left;
    private PriorityQueue<Integer> right;

    public MedianFinder() {
        left = new PriorityQueue<>((o1, o2) -> o2 - o1);
        right = new PriorityQueue<>((o1, o2) -> o1 - o2);
    }
    
    public void addNum(int num) {
        if (left.isEmpty() || left.peek() >= num) {
            left.offer(num);
        } else {
            right.offer(num);
        }

        if (left.size() > right.size() + 1) {
            right.offer(left.poll());
        } else if (right.size() > left.size()) {
            left.offer(right.poll());
        }
    }
    
    public double findMedian() {
        if (left.size() == right.size()) {
            return (left.peek() + right.peek()) / 2.0;
        }
        return left.peek();
    }
}