파이썬 itertools.accumulate() 함수로 접두사 합(Prefix Sum) 배열 구현하기

배열이 주어졌을 때, 파이썬의 itertools.accumulate() 함수를 활용하면 접두사 합(prefix sum), 즉 누적 합 배열을 아주 간단하게 만들 수 있습니다. 이 글에서는 accumulate() 함수의 동작 원리와 함께 실제 예제 코드를 통해 접두사 합 배열을 구하는 방법을 살펴보겠습니다.

accumulate() 함수란?

itertools 모듈의 accumulate(iterable[, func]) 함수는 반복 가능한 객체(iterable)의 요소들을 순서대로 누적 계산하는 반복자(iterator)를 생성해 반환합니다. 기본 동작은 각 위치까지의 누적 합(running total)을 계산하는 것이며, 요소는 Decimal이나 Fraction을 포함해 덧셈이 가능한 모든 타입이 될 수 있습니다.

또한 선택적으로 두 개의 인수를 받는 함수를 전달할 수 있는데, 이 경우 덧셈 대신 해당 함수가 누적 연산에 사용됩니다. 참고로 itertools 모듈의 함수들은 모두 반복자를 생성하고 반환하기 때문에, 스트림을 소진하는 방식(함수 호출이나 루프 등)으로만 접근해야 한다는 점에 유의해야 합니다.

동작 예시

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

위 결과를 단계별로 보면 다음과 같습니다.
첫 번째 요소: 1
두 번째 요소: 1 + 0 = 1
세 번째 요소: 1 + 2 = 3
네 번째 요소: 3 + 3 = 6
다섯 번째 요소: 6 + 5 = 11

알고리즘

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

예제 코드

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

def summation(A):
    print("The List after Summation ::>", list(accumulate(A)))

# 드라이버 프로그램
if __name__ == "__main__":
    A = list()
    n = int(input("Enter the size of the First List ::"))
    print("Enter the Element of First List ::")
    for i in range(int(n)):
        k = int(input(""))
        A.append(k)
summation(A)

실행 결과

Enter the size of the First List ::5
Enter the Element of First List ::
1
2
3
4
5
The List after Summation ::> [1, 3, 6, 10, 15]

위 실행 결과에서 볼 수 있듯이, 사용자가 입력한 5개의 요소 [1, 2, 3, 4, 5]에 대해 accumulate() 함수는 각 위치까지의 누적 합인 [1, 3, 6, 10, 15]를 반환합니다. 이처럼 itertools.accumulate()를 사용하면 별도의 반복문 없이도 접두사 합 배열을 한 줄로 손쉽게 구할 수 있습니다.