무지성 알고리즘 챌린지 카톡방에 들어와서 문제를 푸는 중이다.
문제
널리 잘 알려진 자료구조 중 최대 힙이 있다. 최대 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.
- 배열에 자연수 x를 넣는다.
- 배열에서 가장 큰 값을 출력하고, 그 값을 배열에서 제거한다.
프로그램은 처음에 비어있는 배열에서 시작하게 된다.
입력
첫째 줄에 연산의 개수 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 큰 값을 출력하고 그 값을 배열에서 제거하는 경우이다. 입력되는 자연수는 231보다 작다.
출력
입력에서 0이 주어진 횟수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 큰 값을 출력하라고 한 경우에는 0을 출력하면 된다.
풀이
힙의 경우 우선순위 큐 라는 자료구조를 사용하게 된다.
우선순위 큐를 하나 만들고, 생성 시점에 인자 Comparator.comparingInt(x -> -x)를 통하여 높은 값을 먼저 꺼내도록 만든다.
import java.io.*;
import java.util.*;
public class Main {
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
int n = Integer.parseInt(br.readLine());
PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparingInt(x -> -x));
for (int i = 0; i < n; i++) {
int x = Integer.parseInt(br.readLine());
if (x == 0) {
System.out.println(pq.isEmpty() ? 0 : pq.poll());
continue;
}
pq.add(x);
}
}
}삼항 연산자를 그렇게 좋아하지는 않지만, 위와 같이 한줄에 깔끔하게 정리할 수 있 .
그렇다면 우선순위 큐를 직접 구현한다면 어떨까?
ArrayList로 우선순위 큐 구현하기
현재 우리는, 우선순위 큐에 값을 추가하는 add와 poll 을 구현해야한다.
ArrayList로 구현하게 되면, add와 poll시 swap연산을 Collections.swap을 이용할 수 있다.
코드에 주석으로 설명을 달아 놓았다.
import java.io.*;
import java.util.*;
public class Main {
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
int n = Integer.parseInt(br.readLine());
// 우선순위 큐 초기화
PriorityQueue pq = new PriorityQueue();
for (int i = 0; i < n; i++) {
int x = Integer.parseInt(br.readLine());
if (x == 0) { // x가 0이면
System.out.println(pq.poll()); // 가장 큰 숫자를 꺼내서 출력
} else { // 아니면
pq.add(x); // 우선순위 큐에 숫자 추가
}
}
}
static class PriorityQueue {
private final List<Integer> data;
// 생성자에서 리스트 초기화
PriorityQueue() {
data = new ArrayList<>();
}
// 우선순위 큐에 데이터 추가하는 메소드
public void add(Integer item) {
// item을 리스트에 추가
data.add(item);
int child = data.size() - 1;
// heapify up: 새로 추가된 노드가 잘못된 위치에 있을 경우, 조건을 만족할 때까지 부모 노드와 교환
while (child > 0) {
int parent = (child - 1) / 2;
if (data.get(child) <= data.get(parent)) {
break;
}
Collections.swap(data, child, parent);
child = parent;
}
}
public Integer poll() {
// 문제에서 배열이 비어있는 경우인데 가장 큰 값을 출력하라고 한 경우에는 0을 출력하라고 하여 null이 아닌 0 을 리턴하도록함.
// 큐가 비어있으면 0 반환
if (data.isEmpty()) {
return 0;
}
// 아니면 heapify down 과정 진행
Collections.swap(data, 0, data.size() - 1); // 루트 노드와 마지막 노드 교환
Integer item = data.remove(data.size() - 1); // 마지막 노드 삭제(원래 루트)
int parent = 0;
// heapify down: 루트노드부터 시작해서 자식노드와 비교하여 잘못된 위치에 있으면 자식노드와 교환. 루트에서 아래로 내려감
while (true) {
int leftChild = parent * 2 + 1;
if (leftChild >= data.size()) {
break; // If there are no children, break.
}
int rightChild = leftChild + 1;
int maxChild = leftChild;
// 자식 노드 중에 값이 더 큰 노드를 maxChild로 설정
if (rightChild < data.size() && data.get(rightChild) > data.get(leftChild)) {
maxChild = rightChild;
}
if (data.get(parent) >= data.get(maxChild)) {
break; // 바꿀 필요가 없으면 break
}
Collections.swap(data, parent, maxChild); // 값이 더 큰 자식노드와 바꿈
parent = maxChild;
}
return item;
}
}
}코드가 상당히 길어지게 되었는데, 일단 힙에서는 add에서는 heapify-up과 poll에서는 heapify-down 연산이 수행 되는 것을 알아야된다.
글에서 다루기엔 내용이 너무 길어져서.. 생략한다.
