[Python] 누적합 알고리즘

[1,2,3,4,5,6,7,8,9,10] 리스트가 있다고 했을 때

3번 인덱스부터 마지막 인덱스의 합을 구한다고 하자(인덱스는 1부터 시작)

arr = [1,2,3,4,5,6,7,8,9,10]

start = 3
end = len(arr)
print(sum(arr[start-1:end]))

 

 

 

그렇다면 이 동작을 m 번 반복한다고 해보자

arr = [1,2,3,4,5,6,7,8,9,10]

m = int(input())
for _ in range(m):
    start, end = map(int, input().split())
	print(sum(arr[start-1:end]))

이렇게 되면 O(n*m)의 시간 복잡도가 나온다

 

이렇게 되면 m 이 커질수록 시간 복잡도가 커지는데

이럴 때 누적합을 사용한다

 

DP를 사용하여 누적합을 구한다.

arr = [1,2,3,4,5,6,7,8,9,10]

dp = [0] * (len(arr)+1)

for i in range(1, len(arr)+1):
    dp[i] = dp[i-1] + arr[i-1]
    
print(dp) # [0, 1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

 

 

이 때 A번 인덱스부터 B번 인덱스의 합을 구한다고 하면,

B에서 A-1를 뺀 값을 구하면 된다

 

3번 인덱스부터 마지막 인덱스의 합을 누적합으로 구한다면(인덱스는 1부터 시작)

start = 3
end = len(arr)
print(dp[end] - dp[start-1])

 

이렇게 되었을 때 m번 반복한다면

arr = [1,2,3,4,5,6,7,8,9,10]
dp = [0] * (len(arr)+1)

for i in range(1, len(arr)+1):
    dp[i] = dp[i-1] + arr[i-1]
    
m = int(input())

for _ in range(m):
    start, end = map(int, input().split())
  
    print(dp[end] - dp[start-1])

시간 복잡도는 O(n+m)이 된다.

 

 

 

백준 코테 추천 문제

https://www.acmicpc.net/problem/11659