개똥벌레(3020) / python

acmicpc 코테 준비 문제 - 구현 - 개똥벌레 3020번 - python

문제 설명

제한 사항

• 첫째 줄에 N과 H가 주어진다. N은 항상 짝수이다. (2 ≤ N ≤ 200,000, 2 ≤ H ≤ 500,000)
• 다음 N개 줄에는 장애물의 크기가 순서대로 주어진다. 장애물의 크기는 H보다 작은 양수이다.

문제 풀이

단순하게 생각하면 2차원 배열을 만들어 행 or 열의 합을 계산하고 그 합들 중의 최소값과 그 개수를 구하면 되는 문제지만 최악의 경우 N 200,000 , H 500,000의 크기로 봤을때 2차원 배열을 만들어 처리하면 메모리가 초과하고 2중 for문으로 순회하면 시간 초과가 날 것 처럼 보인다. 필자가 빠르게 적은 다음 코드는 메모리 초과되었다.

import sys
input = sys.stdin.readline
N, H = map(int, input().split(" "))
cave = []

for i in range(N):
    height = int(input())
    if i % 2 == 0:
        temp = [1 if height > j else 0 for j in range(H)]
    else:
        temp = [1 if height > j else 0 for j in range(H)][::-1]
    cave.append(temp)

min_num = min(map(sum, zip(*cave)))
min_cnt = list(map(sum, zip(*cave))).count(min_num)

print(min_num, min_cnt)

누적합

이 문제에서 메모리와 연산시간의 해결책은 누적합(prefix sum)이다.
누적합 S(index)는 배열에서 0 ~ index 까지의 요소들의 합이다.
이를 가지고 문제를 어떻게 해결 하는지 처음에 의문이었는데 주어진 장애물들에 충돌하게 될때 +1을 한다는 점을 생각해보면 특정 높이에 +1을 해놓는 것 만으로도 나중에 누적합을 구했을 때 해결된다는 사실을 알았다.

다음 예시를 보자

H = 7 일때
주어진 장애물의 높이 heigth 를 가지고 prefix_sum 배열의 인덱스 중 하나에 접근해 +1 을
prefix_sum = [0, 0, 0, 0, 0, 0, 0] 에서 heigth 값이 5 ->
prefix_sum = [1, 1, 1, 1, 1, 0, 0] 에서 heigth 값이 3 ->
prefix_sum = [2, 2, 2, 1, 1, 0, 0] 에서 heigth 값이 1 ->
prefix_sum = [3, 2, 2, 1, 1, 0, 0] 이런 식으로 장애물과 충돌하는 횟수의 합을 저장할 수있다.

물론 매번 prefix_sum 배열을 순회하며 1값들을 더해 주는게 아니라 마지막에 전부 누적합한다.

실제
prefix_sum = [0, 0, 0, 0, 0, 0, 0] 에서 heigth 값이 5 ->
prefix_sum = [0, 0, 0, 0, 1, 0, 0] 에서 heigth 값이 3 ->
prefix_sum = [0, 0, 1, 0, 1, 0, 0] 에서 heigth 값이 1 ->
prefix_sum = [1, 0, 1, 0, 1, 0, 0] -> (누적합) -> prefix_sum = [3, 2, 2, 1, 1, 0, 0]

이런 과정에서 연산 속도와 메모리문제를 해결하여 문제를 풀 수 있다. 이 문제에서는 장애물이 위 아래 반복해서 나오기 때문에 이 부분만 고려하여 구현 코드를 작성하면 다음과 같다.

import sys
input = sys.stdin.readline

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

prefix_sum = [0] * h

# 위 아래 장애물을 누적합 배열에 기록
for i in range(n):
    length = int(input())

    if i % 2 == 0:
        prefix_sum[h - length] += 1
    else:
        prefix_sum[0] += 1
        prefix_sum[length] -= 1

# 누적합 계산
for i in range(1, h):
    prefix_sum[i] += prefix_sum[i - 1]

min_num = min(prefix_sum)
min_cnt = prefix_sum.count(min_num)

print(min_num, min_cnt)