Pythonで1次元配列の累積和（ランニングサム）を求める方法

問題の概要

配列 nums が与えられたとします。配列の「ランニングサム（累積和）」とは、rs[i] が nums[0] から nums[i] までのすべての要素の合計を表す新しい配列のことです。最終的には、nums 全体のランニングサムを返します。

例えば、入力が nums = [8,3,6,2,1,4,5] の場合、出力は [8, 11, 17, 19, 20, 24, 29] となります。その理由は以下の通りです。

rs[0] = nums[0]
rs[1] = nums[0..1] の合計 = 8 + 3 = 11
rs[2] = nums[0..2] の合計 = 8 + 3 + 6 = 17
以下同様に続く

解決の手順

この問題を解くには、以下の手順に従います。

• n := nums のサイズ
• rs := [nums[0]]（最初の要素で初期化）
• i を 1 から n-1 まで繰り返す：
  • nums[i] := nums[i] + nums[i-1]
  • rs の末尾に nums[i] を追加
• rs を返す

Pythonでの実装例

理解を深めるために、以下の実装例を見てみましょう。

def solve(nums):
    n = len(nums)
    rs = [nums[0]]

    for i in range(1, n):
        nums[i] += nums[i-1]
        rs.append(nums[i])
    return rs

nums = [8,3,6,2,1,4,5]
print(solve(nums))

入力

[8,3,6,2,1,4,5]

出力

[8, 11, 17, 19, 20, 24, 29]

別のアプローチ：itertools.accumulate を使う

Pythonでは、標準ライブラリの itertools.accumulate を利用すると、同じ結果をより簡潔なコードで得ることができます。

from itertools import accumulate

nums = [8,3,6,2,1,4,5]
print(list(accumulate(nums)))
# 出力: [8, 11, 17, 19, 20, 24, 29]

どちらの方法でも計算量は O(n) となり、配列を一度走査するだけで累積和を求められます。手動で実装する場合はアルゴリズムの仕組みがよく分かり、itertools を使う場合は可読性と簡潔さが向上するというメリットがあります。