🎯

【Java】いもす法

2024/04/12に公開

Java

初心者

tech

いもす法

いもす法とは、0で初期化された配列の指定した区間の要素に同じ数を足した配列を求める時に、累積和を用いることにより高速に求めることができるアルゴリズムです。
また、今までは配列aに対して累積和の配列sを用意していましたが、今回は配列aを累積和の配列に変換することで求めたい配列を用意します。
そして、累積和の配列において要素数0個を考慮しません。

いもす法のアルゴリズム

サイズn+1の配列aを0で初期化。 ※サイズをn+1とするのは、a[r+1] -= xを行うため。
q回にわたって区間l→rにxを足す。
1. q回、a[l]にxを、a[r+1]に-xを加算します。
  - iを0からq-1までループするfor文を定義して、
  - a[l] += x, a[r+1] -= xを加算します。
2. 配列aを累積和の配列にする。
  - iを1からn-1までループするfor文を定義して、
  - a[i] += a[i-1]と更新します。

`a[l]`に`x`を、`a[r+1]`に`-x`を加算する理由

これまでの学習で、累積和の配列sのi+1番目とi番目の差が、元の配列aの要素a[i] であることはわかっています。これを一般化すると下記の通りになります。
```
s[i+1] - s[i] = a[i]
```

次に、累積和の配列sの区間[l+1] → [r+1] にxを加えるためには、
累積和の配列sにおいてl+1番目にxを加えることと、
累積和の配列sにおいてr+2番目にxを加えた効果が打ち消す必要があります。
1. 累積和の配列sにおいてl+1番目にxを加えることは、下記の通りです。
```
s[l+1] - s[l] += x, 
```
この式に先のs[i+1] - s[i] = a[i]を代入すると、下記の通りになります
```
a[l] += x
```
これで a[l]にxを加算 することになりました。
1. 累積和の配列sにおいてr+2番目にxを加えた効果が打ち消すことは、下記の通りです。
```
s[r+2] - s[r+1] -= x
```
この式に先のs[i+1] - s[i] = a[i]を代入すると、下記の通りになります
```
a[r+1] -= x
```
これで a[r+1]に-xを加算 することになりました。

例題

n = 5
q = 3
q1 : l=2, r=3, x=40
q2 : l=0, r=3, x=1
q3 : l=1, r=4, x=5

アルゴリズム1

サイズ6の配列aを0で初期化。
```
a = [0, 0, 0, 0, 0, 0]
```

アルゴリズム2

3回、a[l]にxを、a[r+1]に-xを加算します。
1. l=2, r=3, x=40
```
a = [0, 0, 40, 0, -40, 0]
```
1. l=0, r=3, x=1
```
a = [1, 0, 40, 0, -41, 0]
```
1. l=1, r=4, x=5
```
a = [1, 5, 40, 0, -41, -5]
```
a[i] += a[i-1]を繰り返し、配列aを累積和の配列にする。
```
a = [1, 6, 46, 46, 5]
```

コード例

今回もコード例の解説は割愛します。

import java.util.*;

public class Main {

    static void printArray(int[] a, int n) {
        for (int i = 0; i < n; i++) {
            if (i != 0) {
                System.out.print(" ");
            }
            System.out.print(a[i]);
        }
        System.out.println();
    }

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n, q;
        n = sc.nextInt();
        q = sc.nextInt();

        // サイズ n+1 の配列 a を作成
        int[] a = new int[n + 1];

        // 
        for (int i = 0; i < q; i++) {
            int l, r, x;
            l = sc.nextInt();
            r = sc.nextInt();
            x = sc.nextInt();

            // a[l] に x を、a[r + 1] に -x を加算;
            a[l] += x;
            a[r + 1] -= x;
        }

        
        // i を 1 から n-1 までループする
        for (int i = 1; i < n; i++) {
            // 配列 a を累積和の配列にする
            a[i] += a[i - 1];
        }
        
        printArray(a, n);
    }
}

入力値

出力結果

1 6 46 46 5

例題の解答と同じになることが分かりました。