백준 - 3020번 - 개똥벌레[파이썬(python)]

문제 출처 : https://www.acmicpc.net/problem/3020

3020번: 개똥벌레

 

문제

개똥벌레 한 마리가 장애물(석순과 종유석)로 가득 찬 동굴에 들어갔다. 동굴의 길이는 N미터이고, 높이는 H미터이다. (N은 짝수) 첫 번째 장애물은 항상 석순이고, 그다음에는 종유석과 석순이 번갈아가면서 등장한다.

아래 그림은 길이가 14미터이고 높이가 5미터인 동굴이다. (예제 그림)

이 개똥벌레는 장애물을 피하지 않는다. 자신이 지나갈 구간을 정한 다음 일직선으로 지나가면서 만나는 모든 장애물을 파괴한다.

위의 그림에서 4번째 구간으로 개똥벌레가 날아간다면 파괴해야 하는 장애물의 수는 총 여덟 개다. (4번째 구간은 길이가 3인 석순과 길이가 4인 석순의 중간지점을 말한다)

하지만, 첫 번째 구간이나 다섯 번째 구간으로 날아간다면 개똥벌레는 장애물 일곱 개만 파괴하면 된다.

동굴의 크기와 높이, 모든 장애물의 크기가 주어진다. 이때, 개똥벌레가 파괴해야 하는 장애물의 최솟값과 그러한 구간이 총 몇 개 있는지 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N과 H가 주어진다. N은 항상 짝수이다. (2 ≤ N ≤ 200,000, 2 ≤ H ≤ 500,000)

다음 N개 줄에는 장애물의 크기가 순서대로 주어진다. 장애물의 크기는 H보다 작은 양수이다.

 

출력

첫째 줄에 개똥벌레가 파괴해야 하는 장애물의 최솟값과 그러한 구간의 수를 공백으로 구분하여 출력한다.

 

구현 코드

n, h = map(int, input().split())

down = [0] * (h + 1)  # 석순
up = [0] * (h + 1)  # 종유석
for i in range(n):
    if i%2 == 0:
        down[int(input())] += 1
    else:
        up[int(input())] += 1

for i in range(h-1, 0, -1):
    down[i] += down[i+1]
    up[i] += up[i+1]

min_value = n
min_count = 0

for i in range(1, h+1):
    x = down[i] + up[h-i+1]
    if x < min_value:
        min_count = 1
        min_value = x
    elif x == min_value:
        min_count += 1
print(min_value, min_count)

 

문제 해설

위 문제에서 단순히 동굴의 길이 N 을 높이인 H만큼 완전 탐색한다고 하면 O(NH)의 시간 복잡도를 가지게 된다.

먼저 O(NH)의 시간 복잡도로는 문제를 풀 수 없다. 왜냐하면 N = 10만이고 입력된 모든 석순과 종유석의 길이가 10만이라고 하면 N * H = 100억이 되기 때문이다. 

 

위 문제에서 시간 복잡도를 줄이기 위해 이분 탐색으로도 풀 수 있지만 Prefix Sum을 사용하여 풀 수도 있다.

이 글에선 Prefix Sum을 통한 풀이법을 작성하였다. Prefix Sum을 사용하였을 경우 시간 복잡도는 O(N + 2H)가 된다.

위의 예시의 경우 대략 30만 번의 연산으로 문제를 해결할 수 있다.

 

문제의 쉬운 접근을 위해 석순만 존재한다고 가정해보자.

길이가 1, 2, 3, 4, 5인 석순이 있다고 가정하면, 개똥벌레가 높이 1로 날아갈 경우 5개의 석순에 부딪히게 된다. 2로 날 경우 길이가 1인 석순은 부딪히지 않으므로 4개의 석순에 부딪히게 된다. 리스트 down의 인덱스가 개똥벌레가 날아갈 높이이고 리스트의 값이 개똥벌레가 부딪히는 석순이라고 하면 down = [0, 5, 4, 3, 2, 1](개똥벌레는 0의 높이로 날 수 없으로 첫 번째 값은 0이다.) 이 될 것이다. 그렇다면 이 down 리스트는 Prefix Sum을 활용하여 만들어 줄 수 있다. 아래는 길이가 1, 2, 3, 4, 5인 석순이 있을 때 down 리스트 생성과정이다. 

 

• 먼저 H+1 길이의 0으로 초기화된 리스트를 생성한다 => [0, 0, 0, 0, 0, 0]
• 석순의 길이를 입력받아 길이를 리스트의 인덱스로 갖는 값에 +1을 해준다. => [0, 1, 1, 1, 1, 1]  
• 리스트의 인덱스 H-1부터 시작하여 1번째 인덱스까지 Prefix sum을 구한다 => [0, 5, 4, 3, 2, 1]

 

위와 같은 과정을 거치면 개똥벌레가 날아갈 높이가 i라고 할 때, down[i]를 통해 높이 i일때 부딪히는 석순의 개수를 O(1)의 시간 복잡도로 구할 수 있다.

종유석의 경우 down과 동일하게 Prefix Sum을 활용하여 up리스트를 만들어주면 개똥벌레가 날아갈 높이가 i 라고 할때 up[H-i+1]을 통해 개똥벌레가 부딪히는 종유석의 개수를 구할 수 있다. i를 H만큼 반복하며 최솟값을 찾고 그 개수를 count 하면 답을 구할 수 있다.

Prefix Sum을 구하는데 O(N+H)의 시간 복잡도 높이를 i부터 H까지 순회하는데 O(H)의 시간 복잡도를 가지므로 O(N+2H)의 시간 복잡도로 답을 구할 수 있다.