itertools.accumulate 정리

파이썬으로 PS를 하다 보면 구간합 배열을 만들 일이 자주 생긴다. for 문으로 직접 쓰던 걸 itertools.accumulate 하나로 대체할 수 있다.

기본 사용법

누적 연산 결과를 하나씩 내놓는 이터레이터를 반환한다. 기본 동작은 누적합이다.

from itertools import accumulate
 
list(accumulate([1, 2, 3, 4, 5]))
# [1, 3, 6, 10, 15]

accumulate(iterable, func=operator.add, *, initial=None)

• 결과 개수는 입력과 같다 (initial을 주면 +1)
• 리스트가 아니라 이터레이터라서 필요하면 list()로 감싸야 한다

list(accumulate([1, 2, 3], lambda acc, x: acc * 2 + x))
# [1, 4, 11]   →  1, 1*2+2=4, 4*2+3=11

dp[i] = dp[i-1] * 2 + a[i] 같은 점화식을 한 줄로 굴릴 수 있다.

initial (Python 3.8+)

list(accumulate([1, 2, 3, 4], initial=0))
# [0, 1, 3, 6, 10]

구간합 배열을 만들 때 이게 제일 깔끔하다. [0] + list(accumulate(a))와 결과는 같지만 리스트 연결이 없어서 조금 더 빠르다.

import sys
from itertools import accumulate
 
input = sys.stdin.readline
n, m = map(int, input().split())
a = list(map(int, input().split()))
pre = list(accumulate(a, initial=0))
 
for _ in range(m):
    i, j = map(int, input().split())      # 1-indexed, 닫힌 구간
    print(pre[j] - pre[i - 1])