백준 1461번 도서관

 Description

세준이는 도서관에서 일한다. 도서관의 개방시간이 끝나서 세준이는 사람들이 마구 놓은 책을 다시 가져다 놓아야 한다. 세준이는 현재 0에 있고, 사람들이 마구 놓은 책도 전부 0에 있다. 각 책들의 원래 위치가 주어질 때, 책을 모두 제자리에 놔둘 때 드는 최소 걸음 수를 계산하는 프로그램을 작성하시오. 세준이는 한 걸음에 좌표 1칸씩 가며, 책의 원래 위치는 정수 좌표이다. 책을 모두 제자리에 놔둔 후에는 다시 0으로 돌아올 필요는 없다. 그리고 세준이는 한 번에 최대 M권의 책을 들 수 있다.

 Input

첫째 줄에 책의 개수 N과, 세준이가 한 번에 들 수 있는 책의 개수 M이 주어진다. 둘째 줄에는 책의 위치가 주어진다. N은 10,000보다 작거나 같은 자연수이고, M은 10,000보다 작거나 같다. 책의 위치는 0이 아니며, 그 절댓값이 10,000보다 작거나 같다.

 Output

첫째 줄에 정답을 출력한다.

 Solution

  • 1차 시도: 20%에서 런타임 에러 판정을 받았다.
  • 반례
5 2
1 2 3 4 5
n, m = map(int, input().split())
books = list(map(int, input().split()))

left = []
right = []

for x in books:
    if x < 0:
        left.append(-x)
    else:
        right.append(x)

sum = 0
def go(list):
    global sum
    global m
    while(len(list)>0):
        for i in range(m):
            if len(list) == 0:
                break
            elif i == 0:
                sum += (2 * list.pop())
            else:
                list.pop()

left.sort()
right.sort()

last = 0
if max(left) >= max(right):
    last = max(left)
    go(right)
    go(left)
else:
    last = max(right)
    go(left)
    go(right)

print(sum-last)
  • 2차 시도: left 또는 right 배열이 존재하지 않을 때를 고려하여 조건을 추가하니 정답 판정을 받았다.
n, m = map(int, input().split())
books = list(map(int, input().split()))

left = []
right = []

for x in books:
    if x < 0:
        left.append(-x)
    else:
        right.append(x)

sum = 0
def go(list):
    global sum
    global m
    while(len(list)>0):
        for i in range(m):
            if len(list) == 0:
                break
            elif i == 0:
                sum += (2 * list.pop())
            else:
                list.pop()

left.sort()
right.sort()

last = 0
if len(left) == 0:
    last = max(right)
    go(right)
elif len(right) == 0:
    last = max(left)
    go(left)
elif max(left) >= max(right):
    last = max(left)
    go(right)
    go(left)
else:
    last = max(right)
    go(left)
    go(right)

print(sum-last)

 Feedback

  • 우선순위 큐를 사용하면 조금 더 간단하게 풀 수 있다.
  • 파이썬은 최소 힙(Min Heap)이 구현되어 있기 때문에 이 문제에서처럼 최대 힙(Max Heap) 구현을 위해서는 원소를 음수로 삽입해주어야 한다.
  • 파이썬 heapq 사용법
import heapq

n, m = map(int, input().split(' '))
array = list(map(int, input().split(' ')))
positive = []
negative = []

largest = max(max(array), - min(array))


for i in array:
    if i > 0:
        heapq.heappush(positive, -i)
    else:
        heapq.heappush(negative, i)

result = 0

while positive:
    result += heapq.heappop(positive)
    for _ in range(m - 1):
        if positive:
            heapq.heappop(positive)

while negative:
    result += heapq.heappop(negative)
    for _ in range(m - 1):
        if negative:
            heapq.heappop(negative)

print(-result * 2 - largest)
  • 파이썬 코드도 길길래 자바 코드도 봤는데 역시 어마무시하게 길다.
import java.util.ArrayList;
import java.util.PriorityQueue;
import java.util.Scanner;

public class Main {

	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);

		int n = sc.nextInt();
		int m = sc.nextInt();
		int max_value = -10001;
		int min_value = 10001;

		ArrayList<Integer> data = new ArrayList<Integer>();
		// 자바(Java)에서는 기본적으로 우선순위 큐가 최소 힙(Min Heap)
		PriorityQueue<Integer> positive = new PriorityQueue<Integer>();
		PriorityQueue<Integer> negative = new PriorityQueue<Integer>();

		// 모든 센서의 위치를 입력 받아 오름차순 정렬
		for (int i = 0; i < n; i++) {
			int x = sc.nextInt();
			data.add(x);
			max_value = Math.max(max_value, x);
			min_value = Math.min(min_value, x);
		}

		// 가장 거리가 먼 책까지의 거리
		int largest = Math.max(max_value, -min_value);

		// 최대 힙(Max Heap)을 이용
		for (int i = 0; i < n; i++) {
			// 책의 위치가 양수인 경우
			if (data.get(i) > 0) positive.offer(-data.get(i));
			// 책의 위치가 음수인 경우
			else negative.offer(data.get(i));
		}

		int result = 0;
		while (!positive.isEmpty()) {
			// 한 번에 m개씩 옮길 수 있으므로 m개씩 빼내기
			result += positive.poll();
			for (int i = 0; i < m - 1; i++) {
				if (!positive.isEmpty()) positive.poll();
			}
		}
		while (!negative.isEmpty()) {
			// 한 번에 m개씩 옮길 수 있으므로 m개씩 빼내기
			result += negative.poll();
			for (int i = 0; i < m - 1; i++) {
				if (!negative.isEmpty()) negative.poll();
			}
		}
		// 일반적으로 왕복 거리를 계산하지만, 가장 먼 곳은 편도 거리 계산
		System.out.println(-result * 2 - largest);
	}
}

Source