[백준/Python] 11659 : 구간 합 구하기 4

 

 

 

import sys 

N, M = map(int, input().split())
arr = list(map(int, sys.stdin.readline().split()))
for _ in range(M):
    i, j = map(int, input().split())
    print(sum(arr[i-1:j]))

 

처음에 작성했던 코드이다.

예제로 테스트해본 결과 정답은 올바르게 나오지만, 아무래도 합계를 구하는 sum(arr[i - 1 : j]) 부분의 시간 복잡도가 매우 커지는 경우가 있게 되어 시간초과 오류로 틀렸음

매번 구간 합을 계산하기 위해 sum(arr[i - 1 : j]) 을 호출하여 M 번의 구간 합을 계산해야 하기 대문이다.

 

이 문제를 해결하기 위해 찾아보니 구간 합 배열 (누적 합)을 구현, 즉 구간 합 배열을 미리 계산해두면 각 구간의 합을 O(1) 시간에 구할 수 있어 전체 시간 복잡도가 크게 감소할 수 있다고 한다. 

따라서 구간 합 배열을 적용하여 수정해본 코드는 다음과 같다.

 

import sys
input = sys.stdin.readline

# 입력 처리
N, M = map(int, input().split())
arr = list(map(int, input().split()))

# 누적 합 배열 생성
prefix_sum = [0] * (N + 1)
for idx in range(1, N + 1):
    prefix_sum[idx] = prefix_sum[idx - 1] + arr[idx - 1]

# 구간 합 계산 및 출력
results = []
for _ in range(M):
    i, j = map(int, input().split())
    results.append(prefix_sum[j] - prefix_sum[i - 1])

# 한 번에 출력
print("\n".join(map(str, results)))

 

먼저 arr의 누적 합을 저장할 배열 prefix_sum 을 만들었으며, 이 배열의 각 요소는 arr의 시작부터 해당 인덱스까지의 합을 저장하게 된다.

이후 i부터 j 까지의 구간의 합은 누적 합 배열에서 prefix_sum[j] - prefix_sum[i-1] 를 통해 계산하여 results 리스트에 append 하였음

마지막으로 results 에 들어간 모든 합계들을 join 으로 한 번에 출력하도록 구현하였다.

이렇게 했더니 시간 초과가 발생하지 않고 성공했음!