알고리즘/Heap 힙

[백준] 1655번: 가운데를 말해요 - java

아뵹젼 2023. 2. 13.
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.TreeSet;

public class Main {

	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
		StringBuilder sb = new StringBuilder();
		PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
		PriorityQueue<Integer> minHeap = new PriorityQueue<>();
		
		int n = Integer.parseInt(br.readLine());
		int num;
		for(int i=0; i<n; i++) {
			num = Integer.parseInt(br.readLine());
			
			if ( maxHeap.size() == minHeap.size() ) {
				if (minHeap.size() > 0 && num > minHeap.peek()) {
					maxHeap.add(minHeap.poll());
					minHeap.add(num);
				}else {
					maxHeap.add(num);
				}
			} else {
				if (maxHeap.size() > 0 && num < maxHeap.peek()) {
					minHeap.add(maxHeap.poll());
					maxHeap.add(num);
				} else {
					minHeap.add(num);
				}
			}
			
			sb.append(maxHeap.peek()).append("\n");
		}
		bw.write(sb.toString());
		bw.flush();
		bw.close();
	}

}

 

이 문제는 시간제한이 0.1초로 극악이기 때문에, 단순한 자료구조 탐색으로는 절대 풀 수 없다.

이 문제를 풀 수 있는 다른 방법이 있는지 모르겠으나, 대부분의 사람들이 우선순위 큐를 이용했고, 나도 해당 방법을 택했다.

 

우선순위큐 2개로 각각 최대힙(중간값 이하의 숫자들), 최소힙(중간값 이상의 숫자들)을 만들고,

두 힙이 같은 사이즈를 유지하면서 중간값을 O(1) 로 가져올 수 있도록 하는 문제이다.

 

중간값은 항상 최대힙.peek()으로 가져온다는 것! 을 명심해야 한다.

1 5 2 10 -99 7 5 를 예로 들자.

최대힙 : 5 2 1 -99

최소힙 : 5  7 10

의 형태를 띠게 될 것이고, 중간값은 최대힙의 head 인 5가 된다.

 

이렇게 최대힙, 최소힙을 유지하기 위해선 숫자 (num) 를 삽입할 때 다음과 같은 4가지 조건을 따라야 한다.

 

1. 최대힙과 최소힙의 크기가 같은 경우, num 이 최소힙의 head (중간값 이상의 값) 보다 크다면?

-> num은 중간값 보다 큰 값이므로, 최소힙에 들어가야 한다. 따라서 최소힙의 head(minHeap.peek())를 최대힙에 넣고, 최소힙에 num을 넣는다.

2. 최대힙과 최소힙의 크기가 같은 경우, num 이 최소힙의 head 보다 크지 않다면? 

-> num은 중간값이 되거나, 중간값 이하의 값이 되야 하므로 최대힙에 num을 넣는다.

3. 최대힙의 크기가 최소힙보다 큰 경우, num이 중간값보다 작다면?

-> num은 최대힙에 들어가야 하는데, 두 힙의 크기를 맞추기 위해 중간값(maxHeap.poll())을 최소힙에 넣는다.

4. 최대힙의 크기가 최소힙보다 큰 경우, num이 중간값 이상이라면?

-> num을 최소힙에 넣는다.

 

즉, 정리하자면 다음과 같다.

최대힙에는 중간값 이하의 값들이 저장되어 있으며, 중간값은 최대힙.peek()이다.

최소힙에는 중간값 이상의 값들이 저장되어 있다.

최대힙과 최소힙은 항상 사이즈를 같게 유지하며, 사이즈가 같을 경우에는 최대힙에 원소를 삽입한다.

사이즈를 유지하기 위해, 원소가 최대힙에 삽입되어야 할 때, 삽입될 원소가 중간값보다 큰 경우, 최소힙의 head를 빼서 최대힙에 삽입하고, 해당 원소는 최소힙에 넣는다. -> 원소와 최소힙의 head를 교체!!

사이즈를 유지하기 위해, 원소가 최소힙에 삽입되어야 할 때, 삽입될 원소가 중간값보다 작은 경우, 최대힙의 head(중간값)을 빼서 최소힙에 삽입하고, 해당 원소는 최대힙에 넣는다. -> 원소와 최대힙의 head(중간값) 을 교체!!

댓글