Pythonで範囲合計を効率的に求められるデータ構造を定義する方法

整数のリストから構築できるデータ構造を作り、必要なときにいつでもインデックス i から j−1 までの要素の合計を効率的に求められるようにしたいと考えます。このデータ構造には、次の2つの機能を実装します。

• 整数配列を受け取って新しいインスタンスを初期化するコンストラクタ
• 開始インデックス i から終了インデックス j−1 までの要素の合計を返す get_sum(i, j)

問題の例

たとえば、入力が array = [5,2,3,6,4,7,8,9,3,2] の場合を考えてみましょう。オブジェクト obj を生成し、obj.get_sum(1,5) と obj.get_sum(4,8) を呼び出すと、結果はそれぞれ 15 と 28 になります。これは、1つ目の範囲の要素が [2,3,6,4] でその合計が 15、2つ目の範囲の要素が [4,7,8,9] でその合計が 28 だからです。

解決のアプローチ：累積和（プレフィックスサム）

この問題は「累積和」というテクニックを使えば、各クエリを高速に処理できます。具体的な手順は以下の通りです。

• コンストラクタを定義する。引数として配列を受け取る
• sums := リストを用意し、初期値として 0 を挿入する
• 配列内の各要素 x に対して、次の処理を行う
  • sums の末尾に「x + sums の最後の要素」を追加する
• get_sum() 関数を定義する。引数として i と j を受け取る
• sums[j] − sums[i] を返す

こうすることで、sums[k] には「先頭から k 番目までの要素の合計」が格納されるため、任意の区間の合計は2つの累積和の差として即座に計算できるのです。

実装例

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

class RangeSum:
    def __init__(self, array):
        self.sums = [0]
        for x in array:
            self.sums.append(x + self.sums[-1])
    def get_sum(self, i, j):
        return self.sums[j] - self.sums[i]

array = [5,2,3,6,4,7,8,9,3,2]
obj = RangeSum(array)
print(obj.get_sum(1,5))
print(obj.get_sum(4,8))

入力

[5,2,3,6,4,7,8,9,3,2]
obj.get_sum(1,5)
obj.get_sum(4,8)

出力

15
28

計算量のポイント

コンストラクタでは配列を一度だけ走査するため O(n) の時間がかかりますが、その後の get_sum の呼び出しはすべて O(1) で完了します。同じ配列に対して何度も範囲合計を求めるようなケースでは、クエリごとにループで合計を計算する方法（1回あたり O(n)）と比べて、大幅に高速に動作します。