[Prefix Sum] Python - 백준 3020번: 개똥벌레

www.acmicpc.net/problem/3020

import sys
input = sys.stdin.readline

N, H = map(int, input().split())
bottomObstacles = [0]*H
topObstacles = [0]*H

for i in range(N//2):
    o1, o2 = int(input()), int(input())
    bottomObstacles[o1-1] += 1
    topObstacles[o2-1] += 1

bottomPrefixSum = [0]*(H+1)
topPrefixSum = [0]*(H+1)
for i in range(H-1, -1, -1):
    bottomPrefixSum[i] = bottomPrefixSum[i+1] + bottomObstacles[i]
    topPrefixSum[i] = topPrefixSum[i+1] + topObstacles[i]

minObstaclesToBeDestroyed, minCount = 400001, 0
for b, t in zip(range(H), range(H-1, -1, -1)):
    cur = bottomPrefixSum[b] + topPrefixSum[t]

    if cur == minObstaclesToBeDestroyed:
        minCount += 1
    elif cur < minObstaclesToBeDestroyed:
        minObstaclesToBeDestroyed = cur
        minCount = 1

print(minObstaclesToBeDestroyed, minCount)

이번 풀이에서는 prefix sum 알고리즘이 핵심이다.

우선, bottomObstacles, topObstacles으로 나누어서 해당 높이에 해당하는 장애물의 갯수를 따로 기록해준다.

 

bottomPrefixSum[i]에는 해당 높이(i: 0~H-1)에서 충돌하게 될 장애물들의 갯수를 기록하는데, 이것은 for문을 H-1에서 0까지 역순으로 돌면서 prefix를 구해주면 된다. topPrefixSum은 역으로 저장한다.

 

두 prefix sum의 값을 합치면 해당 위치에서의 장애물의 갯수가 된다. 이 때, topPrefixSum은 역순으로 저장되어 있으므로, for문의 인덱스값을 조절해준다. 아무튼 그렇게 장애물 갯수들의 최소값을 구해주고, 최소값과 현재 i 높이에서의 장애물 갯수가 같을 경우에는 minCount를 증가시켜준다.