[python] 백준 11659 - 구간 합 구하기 4

백준 11659 - 구간 합 구하기4 [SILVER 3]

이름 그대로 누적 합 구하기 문제다.

print(sum(numbers[start : end]))로는 절대 풀 수 없는 문제다.

숫자 배열이 10만, 계산 및 출력해야할 개수가 10만이다.

또한 sum(numbers[0 : 100,000])을 10만번 반복한다고 하면,

100,000 * 100,000. 즉, 최악의 경우 100억 번(n²)의 연산이 필요한 것이다…

물론 딕셔너리나 리스트로 캐시를 만들 수 있겠지만,

그것을 감안하더라도 절대 n log n안으로는 줄일 수 없을 것이다.

누적 합

만약, numbers라는 배열에서, 3번째 수부터 5번째 수까지 합을 구한다고 하자.

이는 1번부터 5번까지의 합에서 1번부터 2번까지의 합을 뺀 것과 마찬가지다.

sum(numbers[2:4]) == sum(numbers[0:4]) - sum(numbers[0:1])

위의 코드로 보면 당연히 전자가 빠르겠지만, 후자는 다른 방법으로 훨씬 쉽게 구할 수 있다.

numbers의 1번째부터 n번째까지 더한 값을, prefix_sum 배열의 n번째 인자로 넣자.

만약 numbers 배열이 [1, 3, 5, 7, 10]이라면,

prefix_sum 배열은 [1, 4, 9, 16, 26]이 된다.

이제 위의 코드를 수정해보자.

sum(numbers[2:4]) == prefix_sum[4] - prefix_sum[1]

prefix_sum 배열을 만들어놨다면, 후자는 1번의 연산으로 정답을 도출해낸다.

prefix_sum, 누적합 배열 만들기

for i in range(N):
    prefix_sum[i] = sum(numbers[0:i])

위의 코드는 이 문제에서 시간 초과로 이어진다.

이전에 말한 최악의 경우일 때, 위의 배열을 채우는 데에만 약 100,000 * 50,000번의 연산이 필요하다.

현재 i번째 값을 채워넣을 때,

prefix_sum 배열의 i-1번 째 값, 즉 이전 루프에서 구한 값과 numbers[i]을 더해주면 된다.

최악의 경우에도 10만번의 연산이 끝이다.

# 연산이 1부터 시작할 때를 위해 prefix_sum[0]에는 0이 들어간다.
prefix_sum[1] = numbers[1]

for i in range(2, N):
    prefix_sum[i] = prefix_sum[i-1] + numbers[i]

정답 코드

# input이 매우 많은 문제이므로 반드시 readline으로 값을 받자

import sys

input = sys.stdin.readline

N, M = map(int, input().split())
numbers = list(map(int, input().split()))
prefix_sum = [0] * (N+1)
prefix_sum[1] = numbers[0]

for i in range(2, N+1):
    prefix_sum[i] = prefix_sum[i-1]+numbers[i-1]
    
for _ in range(M):
    s, e = map(int, input().split())
    print(prefix_sum[e] - prefix_sum[s-1])

결과

메모리	시간
41116KB	248ms