【いもす法】区間に等差数列を加算したい話

この記事でやること
注意事項
いもす法について
発想：2回累積和を取る
初項が0の場合
初項が一般の場合
ソースコード
例題
- 略解

この記事でやること

数列の区間に等差数列を加算するクエリを高速に処理する方法を解説します。

具体的には、数列 $a_0, a_1, \dots, a_{N-1}$ を $0$ で初期化し、以下のクエリを $Q$ 回処理した後の数列を取得することを、全体 $O(N+Q)$ で行います:

$\text{add_line}(l, r, a, d): 各i \in [l, r) に対し, a_i \leftarrow a_i + (a+d(i-l))$

注意事項

ソースコードはPython3.8を用いております。コピペはご自由にどうぞ。（AtCoder等で使用する際はPyPy3で提出して問題ありません。）

いもす法について

おなじみのやつですが、考え方が非常に重要なので確認しておきます。

いもす法は区間全体に同じ数字を加算をします。このとき、各クエリでは「最後に累積和を取ることで元の数列が復元できるような、情報量を削減した数列」を作っていきます。

ということは、同じように「最後に累積和をとることで等差数列を加算した数列が復元できるような、情報量を削減した数列」を作ることができればよさそうです。

発想：2回累積和を取る

先ほどの画像の数列に対し、もう1度累積和を取ってみます。

等差数列が出てきましたね。これを逆に言い換えれば、「2度累積和を取ることで等差数列が復元できるような、情報量を削減した数列」が存在しそうな気がしてきませんか？

というわけで、実際にそのような数列を元の等差数列から構築してみましょう。

初項が0の場合

まずは初項が0の場合を考えます。これは簡単で、下記の画像の一番上の行が作りたかった数列となります。

具体的な処理を一般的に書くと、 $\text{add_line}(l, r, 0, d)$ に対しては、

$a_{l+1} \leftarrow a_{l+1}+d$
$a_r \leftarrow a_{r}-d(r-l)$
$a_{r+1} \leftarrow a_{r+1}+d(r-l-1)$

とすればよいです。

初項が一般の場合

「初項が0の場合」に、初項のズレを加えていい感じに調整してあげればよいです。

具体的な処理を一般的に書くと、 $\text{add_line}(l, r, a, d)$ に対しては、

$a_{l} \leftarrow a_{l}+a$
$a_{l+1} \leftarrow a_{l+1}+d-a$
$a_r \leftarrow a_{r}-d(r-l)-a$
$a_{r+1} \leftarrow a_{r+1}+d(r-l-1)+a$

とすればよいです。

ソースコード

class Imos:
    def __init__(self, N) -> None:
        """
            N := 配列のsize
        """
        self.N = N
        self.A = [0]*(N+2)
    def add_line(self, l, r, a, d) -> None:
        """
            [l, r) に初項a, 公差d の等差数列を加算
        """
        self.A[l] += a
        self.A[l+1] += d-a
        self.A[r] -= d*(r-l)+a
        self.A[r+1] += d*(r-l-1)+a
    def get_array(self) -> None:
        """
            現在の数列を取得
        """
        ret = [0]*self.N
        ret[0] = self.A[0]
        for i in range(self.N-1):
            ret[i+1] = ret[i]+self.A[i+1]
        for i in range(self.N-1):
            ret[i+1] += ret[i]
        return ret