알고리즘

[알고리즘] 누적합 - Python

샥쿠 2023. 8. 7. 22:23

누적합, 구간합

배열의 구간 합을 구할 때 해당 구간을 순회하면서 구하면 최대 O(n)의 시간이 걸리지만, 누적합 배열을 미리 만들어 놓으면 O(1)의 시간 복잡도로 구할 수 있다.

index   0  1  2  3  4
nums = [1, 2, 3, 4, 5]
acc  = [0, 1, 3, 6, 10, 15]

nums[i]부터 nums[j]까지의 합 = acc[j+1] - acc[i]

누적합 배열 acc[i+1]는 nums 배열의 0번 값부터 i번 값까지의 합을 의미한다.

따라서 nums 배열의 i번째 값부터 j번째 값까지의 합은 0번 ~ j번 값까지의 합인 acc[j+1]에서 0번 ~ i-1번 값까지의 합인 acc[i]를 뺀 값이 된다.

누적합과 관련한 백준 문제

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

2559번: Python 풀이

import sys

input = sys.stdin.readline

N, K = map(int, input().split())
nums = list(map(int, input().split()))

acc = [0]
for i in range(N):
    acc.append(acc[i] + nums[i])

max_sum = float("-inf")
for i in range(0, N - K + 1):
    # nums[i] 부터 nums[i + K - 1]까지의 합
    temp_sum = acc[i + K] - acc[i]
    max_sum = max(max_sum, temp_sum)

print(max_sum)