파이썬 accumulate 함수로 접두사 합 배열 쉽게 구현하기

배열이 주어졌을 때 접두사 합(prefix sum) 배열은 첫 번째 요소부터 각 위치까지의 누적 합을 순서대로 담은 배열입니다. 파이썬에서는 itertools.accumulate() 함수를 활용하면 반복문을 직접 작성하지 않고도 아주 간결하게 구현할 수 있습니다.

itertools.accumulate(iterable[, func])는 반복 가능한 객체를 입력받아 누적 결과를 하나씩 반환하는 이터레이터를 생성합니다. 모듈의 함수들은 이터레이터를 만들어 반환하므로, list()나 for 루프처럼 스트림을 소진하는 방식으로만 값에 접근해야 합니다. 기본 동작은 누적 합계이며, 요소는 Decimal이나 Fraction처럼 덧셈이 가능한 어떤 타입이든 사용할 수 있습니다. 또한 두 개의 인자를 받는 함수를 선택적으로 전달하면, 덧셈 대신 그 함수가 누적 연산에 사용됩니다.

예제

입력:
data = [1, 0, 2, 3, 5]
>>> list(accumulate(data))     # 누적 합계 계산
출력:
[1, 1, 3, 6, 11]

알고리즘

1단계: 리스트를 생성한다.
2단계: list(accumulate()) 함수를 사용하여 누적 합계를 계산한다.
3단계: 계산된 전체 합계 결과를 화면에 출력한다.

예제 코드

# accumulate 함수를 사용해
# 접두사 합 배열을 출력하는 파이썬 프로그램
from itertools import accumulate

def summation(A):
    print("합계 처리 후의 리스트 ::>", list(accumulate(A)))

# 메인 프로그램
if __name__ == "__main__":
    A = list()
    n = int(input("첫 번째 리스트의 크기를 입력하세요 ::"))
    print("첫 번째 리스트의 요소를 입력하세요 ::")
    for i in range(int(n)):
        k = int(input(""))
        A.append(k)
    summation(A)

실행 결과

첫 번째 리스트의 크기를 입력하세요 ::5
첫 번째 리스트의 요소를 입력하세요 ::
1
2
3
4
5
합계 처리 후의 리스트 ::> [1, 3, 6, 10, 15]

실행 결과를 보면 각 위치의 값이 다음과 같이 계산됩니다. 첫 번째 값은 1, 두 번째 값은 1+2=3, 세 번째 값은 1+2+3=6, 네 번째 값은 1+2+3+4=10, 마지막 값은 1+2+3+4+5=15입니다. 이처럼 accumulate()를 사용하면 접두사 합 배열을 한 줄의 코드로 손쉽게 얻을 수 있으며, 가독성과 성능 면에서도 직접 구현하는 방식보다 유리합니다.