C++で範囲加算を効率的に解く方法 ― 差分配列（いもす法）の活用

問題概要

サイズnの整数配列が与えられ、すべての要素が0で初期化されているとします。さらに値kが与えられ、k回の更新操作を行います。各操作は [startIndex, endIndex, inc] という3つ組で表され、部分配列 A[startIndex ... endIndex]（startIndexとendIndexを含む）の各要素に inc を加算します。k回の操作をすべて実行した後の、変更後の配列を求めるのがこの問題です。

例えば、入力が length = 5、updates = [[1,3,2],[2,4,3],[0,2,-2]] の場合、出力は [-2, 0, 3, 5, 3] となります。

操作の適用順序（例の確認）

• [1,3,2] を適用 → [0, 2, 2, 2, 0]
• [2,4,3] を適用 → [0, 2, 5, 5, 3]
• [0,2,-2] を適用 → [-2, 0, 3, 5, 3]

解法のアプローチ：差分配列（いもす法）

各操作で範囲内の要素を1つずつ直接更新すると、最悪で O(n × k) の計算量が必要になり非効率です。そこで「差分配列」と呼ばれるテクニック（日本では「いもす法」としても知られています）を活用します。各操作を区間の始点と終点+1の2箇所にだけ記録し、最後に累積和を一度計算するだけで、全体を O(n + k) で求められます。

具体的な手順は以下の通りです。

• サイズnの配列 ret を0で初期化して定義する
• 各更新操作 a[i] について、以下を繰り返す
  • l := a[i][0]（開始インデックス）
  • r := a[i][1] + 1（終了インデックス + 1）
  • ret[l] に a[i][2] を加算する
  • r < n の場合は、ret[r] から a[i][2] を減算する
• i = 1 から n - 1 まで、ret[i] += ret[i - 1] として累積和を計算する
• ret を返す

なぜこの方法が機能するのか

区間 [l, r] への inc の加算は、差分配列上では「位置lでincを加算し、位置r+1でincを減算する」として記録されます。その後、先頭から累積和を計算すると、位置lからrまでは累積値にincが反映され、r+1以降は減算分と相殺されて元に戻ります。その結果、目的の区間だけがincだけ増加した配列が得られる仕組みです。

C++での実装例

以下の実装を見ると、より理解が深まるでしょう。

#include <bits/stdc++.h>
using namespace std;

template <typename T>
void print_vector(vector<T> v){
   cout << "[";
   for(int i = 0; i < v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]" << endl;
}

class Solution {
public:
   vector<int> getModifiedArray(int n, vector<vector<int>> &a) {
      vector<int> ret(n);
      for (int i = 0; i < a.size(); i++) {
         int l = a[i][0];
         int r = a[i][1] + 1;
         ret[l] += a[i][2];
         if (r < n) {
            ret[r] -= a[i][2];
         }
      }
      for (int i = 1; i < n; i++) {
         ret[i] += ret[i - 1];
      }
      return ret;
   }
};

int main(){
   Solution ob;
   vector<vector<int>> v = {{1,3,2},{2,4,3},{0,2,-2}};
   print_vector(ob.getModifiedArray(5, v));
}

入力

5, {{1,3,2},{2,4,3},{0,2,-2}}

出力

[-2, 0, 3, 5, 3]