[Java+Python]11286번, 절댓값 힙, Comparator

목차

 

문제

 

절댓값 힙은 다음과 같은 연산을 지원하는 자료구조이다.

1. 배열에 정수 x (x ≠ 0)를 넣는다.
2. 배열에서 절댓값이 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다.
   절댓값이 가장 작은 값이 여러 개일 때는, 가장 작은 수를 출력하고, 그 값을 배열에서 제거한다.

프로그램은 처음에 비어있는 배열에서 시작하게 된다.

 

입력

 

첫째 줄에 연산의 개수 N(1≤N≤100,000)이 주어진다. 

다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 

만약 x가 0이 아니라면 배열에 x라는 값을 넣는(추가하는) 연산이고, 

x가 0이라면 배열에서 절댓값이 가장 작은 값을 출력하고 그 값을 배열에서 제거하는 경우이다. 

입력되는 정수는 $-2^{31}$보다 크고, $2^{31}$보다 작다.

 

출력

 

입력에서 0이 주어진 횟수만큼 답을 출력한다. 

만약 배열이 비어 있는 경우인데 절댓값이 가장 작은 값을 출력하라고 한 경우에는 0을 출력하면 된다.

 

풀이

 

우선순위 큐 단계는 사실 별 특별할 것 없이 사용법을 연습하는 문제들이라 딱히 올릴 생각이 없었다.

 

하지만 이번 문제를 풀면서 자바의 Comparator에 대한 리마인드와, 파이썬에서 절댓값을 큐에 넣을 때

 

어떻게 처리하는지 배우고 또 흥미로워서, 글을 남겨두기로 했다.

 

먼저 자바의 경우 문제의 조건에 맞는 수를 힙에 넣을 때, 다음과 같이 표현할 수 있다.

PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparingInt((Integer a) -> Math.abs(a)).thenComparingInt(a -> a));

Comparator에 두 가지 비교 기준이 들어간 것을 확인할 수 있는데,

 

이는 첫 번째로 절댓값을 비교해 작은 순서로 정렬하고, 절댓값이 같다면 thenComparingInt()를 사용해

 

원래의 값에 따라 다시 오름차순으로 정렬하는 식이다.

 

참고로 우선순위 큐는 기본이 오름차순이기 때문에, 최대 힙이나 다른 방식으로 사용하려면 다소 처리가 필요하다.

 

이어서 파이썬의 경우이다.

heapq.heappush(heap, (abs(x), x))

여기서는 힙에 튜플을 생성해 추가하고 있는데, 각 요소는 abs(x)와 x로 이루어져 있다.

 

이와 같이 설정하면 첫 번째로는 절댓값을 기준으로, 절댓값이 같다면 원래의 값에 따라 오름차순으로 정렬된다.

 

이후에 출력을 할 때는 해당 힙에 들어있는 튜플의 1번째 인덱스 값을 고르면 되는데,

 

역시 자바에 비해 굉장히 짧고, 흥미로웠다.

 

계속해서 구현 전체를 보자.

 

Java

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Comparator;
import java.util.PriorityQueue;

public class Prob11286 {

	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((Integer a) -> Math.abs(a)).thenComparingInt(a -> a));

		StringBuilder sb = new StringBuilder();

		for (int i = 0; i < n; i++) { 
			int x = Integer.parseInt(br.readLine());

			if (x == 0) {
				if (pq.isEmpty()) {
					sb.append(0).append('\n');
				} else {
					sb.append(pq.poll()).append('\n');
				}
			} else {
				pq.offer(x);
			}
		}

		System.out.println(sb);
	}
}

 

Python

 

import sys
import heapq

n = int(sys.stdin.readline().rstrip())
heap = []

for _ in range(n):
    x = int(sys.stdin.readline().rstrip())

    if x == 0:
        if len(heap) == 0:
            print(0)
        else:
            print(heapq.heappop(heap)[1])
    else:
        heapq.heappush(heap, (abs(x), x))

 

Performance

 

코드 길이와 시간 복잡도가 절반도 되지 않는다!