いもす法の解説記事はいろいろありますが、（僕にとってしっくりくる形で）数式を使って表現したものがあまりなかったので、まとめようと思います。
考えたい問題
次のような問題を考えます。
【問題】
長さN の数列a_{1}, ..., a_{N} が与えられる。この数列に対し、次の操作をQ回行う：
k 回目の操作では、a_{l_{k}}, ..., a_{r_{k}} にx_{k}を加算する。
操作後のa_{1}, ..., a_{N}を求めよ。
【制約】
1\leq N \leq 10^5
1 \leq Q \leq 10^5
1 \leq l_{k} < r_{k} \leq N \ (1\leq k \leq Q)
1 \leq x_{k} \leq 100 (1 \leq k \leq Q)
例
以下 0\leq k \leq Qに対して、k回目の操作を行ったあとのa_{i}をa^{(k)}_{i}と書くことにします。つまり、
\begin{align}
\begin{split}
a^{(0)}_{i} &= a_{i} \ (1\leq i \leq N), \\
a^{(k)}_{i} &= \left\{
\begin{array}{ll}
a^{(k-1)}_{i}+x_{k} & (l_{k}\leq i \leq r_{k}) \\
a^{(k-1)}_{i} & (1\leq i < l_{k} \text{ or } r_{k}<i\leq N)
\end{array}
\right. \ (1\leq k \leq Q)
\end{split}
\end{align}
とします。求めたいものはa^{(Q)}_{1}, ..., a^{(Q)}_{N}です。
N = 7, Q = 3 の場合の数列の変化の例です。太字の部分が更新された項を表します。
| k |
l_{k} |
r_{k} |
x_{k} |
a^{(k)}_{1} |
a^{(k)}_{2} |
a^{(k)}_{3} |
a^{(k)}_{4} |
a^{(k)}_{5} |
a^{(k)}_{6} |
a^{(k)}_{7} |
| 0 |
- |
- |
- |
0 |
0 |
0 |
0 |
0 |
0 |
0 |
| 1 |
1 |
3 |
4 |
4 |
4 |
4 |
0 |
0 |
0 |
0 |
| 2 |
2 |
5 |
1 |
4 |
5 |
5 |
1 |
1 |
0 |
0 |
| 3 |
3 |
7 |
10 |
4 |
5 |
15 |
11 |
11 |
10 |
10 |
愚直な計算量
各操作を愚直に行うと最悪でO(N)の計算量がかかるので、全体での計算量は最悪でO(NQ)です。そのため今回の制約だとO(10^{10})程度の計算量となり、時間がかかります(AtCoderだと１秒当たり10^8回の計算が目安とされる)。
いもす法のアルゴリズム
各操作に対する計算量を落とすことを考えます。
各kに対してa^{(k)}_{0}=0, \ b^{(k)}_{i} := a^{(k)}_{i} - a^{(k)}_{i-1} \ (1\leq i \leq N)とします。
すると、(1)式より
b^{(k)}_{i} = \left\{
\begin{array}{ll}
b^{(k-1)}_{i}+x & (i = l_{k}) \\
b^{(k-1)}_{i}-x & (i = r_{k}+1) \\
b^{(k-1)}_{i} & (i \ne l_{k}, r_{k}+1)
\end{array}
\right.
となっています(r_{k} = Nのときは、最後の場合は存在しないことに注意が必要です）。これより、「a^{(k-1)}_{l_{k}}, ..., a^{(k-1)}_{r_{k}}にx_{k}を足す」という操作は、
- 「b^{(k-1)}_{l_{k}}にx_{k}を足し、b^{(k-1)}_{r_{k}+1}に-x_{k}を足す」（r_{k}=Nの場合は、「b_{l_{k}}にx_{k}を足す」だけでよい）
という操作と同値だと分かります。そしてこの操作はO(1)で行えます。よって、全体として操作にかかる時間が、O(NQ)からO(Q)となりました。
このままだとb^{(Q)}_{1}, ..., b^{(Q)}_{N}を求めただけなので、最後にa^{(Q)}_{1}, ..., a^{(Q)}_{n}の情報を復元する必要があります。今、
\begin{align}
\begin{split}
a^{(k)}_{1} &= b^{(k)}_{1}, \\
a^{(k)}_{2} &= (a^{(k)}_{2} - a^{(k)}_{1}) + a^{(k)}_{1} = b^{(k)}_{2} + a^{(k)}_{1}, \\
&..., \\
a^{(k)}_{n} &= (a^{(k)}_{n} - a^{(k)}_{n-1}) + a^{(k)}_{n-1} = b^{(k)}_{n} + a^{(k)}_{n-1}
\end{split}
\end{align}
となるので、(2)式より数列a^{(Q)}_{1}, ..., a^{(Q)}_{N}はb^{(Q)}_{1}, ..., b^{(Q)}_{N}の累積和として得られます。従ってa^{(Q)}_{1}, ..., a^{(Q)}_{N}を求めるのにかかる時間はO(N)です。
よって全体ではO(Q+N)で計算できるようになりました。
例の再考
最初に考えた例の場合、a^{(Q)}_{i}とb^{(Q)}_{i}は次のように更新されます。太字の部分が更新された項を表します。
| k |
l_{k} |
r_{k} |
x_{k} |
a^{(k)}_{1} |
a^{(k)}_{2} |
a^{(k)}_{3} |
a^{(k)}_{4} |
a^{(k)}_{5} |
a^{(k)}_{6} |
a^{(k)}_{7} |
| 0 |
- |
- |
- |
0 |
0 |
0 |
0 |
0 |
0 |
0 |
| 1 |
1 |
3 |
4 |
4 |
4 |
4 |
0 |
0 |
0 |
0 |
| 2 |
2 |
5 |
1 |
4 |
5 |
5 |
1 |
1 |
0 |
0 |
| 3 |
3 |
7 |
10 |
4 |
5 |
15 |
11 |
11 |
10 |
10 |
| k |
l_{k} |
r_{k} |
x_{k} |
b^{(k)}_{1} |
b^{(k)}_{2} |
b^{(k)}_{3} |
b^{(k)}_{4} |
b^{(k)}_{5} |
b^{(k)}_{6} |
b^{(k)}_{7} |
| 0 |
- |
- |
- |
0 |
0 |
0 |
0 |
0 |
0 |
0 |
| 1 |
1 |
3 |
4 |
4 |
0 |
0 |
-4 |
0 |
0 |
0 |
| 2 |
2 |
5 |
1 |
4 |
1 |
0 |
-4 |
0 |
-1 |
0 |
| 3 |
3 |
7 |
10 |
4 |
1 |
10 |
-4 |
0 |
-1 |
0 |
２次元への拡張
いもす法のいいところは、次元を上げても適用できることです。次のような問題を考えます。
【問題】
H\times W 行列A = a_{i, j} が与えられる。次の操作をQ回行う：k回目の操作では、u_{k}\leq i \leq d_{k}かつl_{k}\leq j \leq r_{k}なる(i, j) に対して、a_{i, j}にx_{k}を加算する。
操作後の行列A を求めよ。
【制約】
1\leq HW \leq 10^5
1 \leq Q \leq 10^5
1 \leq l_{i} < r_{i} \leq N \ (1\leq i \leq Q)
1 \leq x_{i} \leq 100 (1 \leq i \leq Q)
この問題も、各操作で愚直にa_{i, j}たちを更新していると最悪計算量がO(QHW)となり時間がかかります。1次元の場合と同様に、階差数列を考えて、各操作をO(1)で行うことを考えます。
この問題を考える上では、操作を一度も行っていない状態のAの各項は全て0としてよいです。0でない項がある状態から始まるときは、各項が0のH\times W 行列A'に対して操作を行い、操作終了後のA'をAに足せば、求めたい行列が求まるからです。
以下、i=0 またはj=0のとき、a_{i, j} = 0とします。k 回目の操作を行った後のa_{i, j}のことをa^{(k)}_{i, j}と置きます。つまり、
\begin{align}
\begin{split}
a^{(0)}_{i, j} &= a_{i, j} = 0 \ (1\leq i \leq H, \ 1\leq j \leq W), \\
a^{(k)}_{i, j} &= \left\{
\begin{array}{ll}
a^{(k-1)}_{i, j} + x_{k} & (u_{k} \leq i \leq d_{k} \text{ and } l_{k} \leq j \leq r_{k}) \\
a^{(k-1)}_{i} & (\text{otherwise})
\end{array}
\right. \ (1\leq k \leq Q)
\end{split}
\end{align}
とします。また、各0\leq k \leq Q, \ 1\leq i \leq H, 1 \leq j \leq Wについて、b^{(k)}_{i, j} := a^{(k)}_{i, j} - a^{(k)}_{i-1, j}, \ c^{(k)}_{i, j} := a^{(k)}_{i, j} - a^{(k)}_{i, j-1}と置きます。 すると(3)より、1\leq k \leq Qに対して
\begin{align}
b^{(k)}_{i, j} &= \left\{
\begin{array}{ll}
b^{(k-1)}_{i, j}+x_{k} & (i = u_{k} \text{ and } l_{k}\leq j \leq r_{k}) \\
b^{(k-1)}_{i, j}-x_{k} & (i = d_{k} + 1 \text{ and } l_{k}\leq j \leq r_{k}) \\
b^{(k-1)}_{i, j} & (\text{otherwise})
\end{array}
\right. , \\
c^{(k)}_{i, j} &= \left\{
\begin{array}{ll}
c^{(k-1)}_{i, j}+x_{k} & (u_{k} \leq i \leq d_{k} \text{ and } j = l_{k}) \\
c^{(k-1)}_{i, j}-x_{k} & (u_{k} \leq i \leq d_{k} \text{ and } j = r_{k} + 1) \\
c^{(k-1)}_{i, j} & (\text{otherwise})
\end{array}
\right.
\end{align}
が成り立ちます。これより、「u_{k}\leq i \leq d_{k}かつl_{k}\leq j \leq r_{k}なる(i, j) に対して、a^{(k-1)}_{i, j}にx_{k}を加算する」という操作は、
-
i = u_{k} かつ l_{k}\leq j \leq r_{k} なる(i, j)について、b^{(k-1)}_{i, j}にx_{k}を加算する。
-
i = d_{k}+1 かつ l_{k}\leq j \leq r_{k} なる(i, j)について、b^{(k-1)}_{i, j}に-x_{k}を加算する。
-
u_{k} \leq i \leq d_{k} かつ j=l_{k} なる(i, j)について、c^{(k-1)}_{i, j}にx_{k}を加算する。
-
u_{k} \leq i \leq d_{k} かつ j = r_{k}+1 なる(i, j)について、c^{(k-1)}_{i, j}に-x_{k}を加算する。
という４つの操作を行うことと同値です。よって各iに対して\{b^{(k)}_{i, j}\}_{j}にいもす法を適用し、各j に対して\{c^{(k)}_{i, j}\}_{i}にいもす法を適用すれば、全体でO(Q + HW)の計算量でb^{(Q)}_{i, j}, \ c^{(Q)}_{i, j}を求められることになります。
しかし、実際はもう少し簡単にb^{(Q)}_{i, j}, \ c^{(Q)}_{i, j}を求められます。
1\leq i\leq H, \ 1\leq j \leq Wに対して、e^{(0)}_{i, j} = 0とします。また、1\leq k \leq Qに対して
e^{(k)}_{i, j} =
\left\{
\begin{array}{ll}
e^{(k-1)}_{i, j}+x_{k} & ((i, j) = (u_{k}, l_{k}), \ (d_{k}+1, r_{k}+1)) \\
e^{(k-1)}_{i, j}-x_{k} & ((i, j) = (u_{k}, r_{k}+1), \ (d_{k}+1, l_{k})) \\
e^{(k-1)}_{i, j} & (\text{otherwise})
\end{array}
\right.
と置きます。さらに、i=0またはj=0のとき、b_{i, j} = c_{i, j} = 0 とします。このとき、b^{(k)}_{i, j} - b^{(k)}_{i, j-1} = e^{(k)}_{i, j}かつc^{(k)}_{i, j} - c^{(k)}_{i-1, j} = e^{(k)}_{i, j}が、すべての1\leq k \leq Q, \ 1\leq i \leq H, \ 1\leq j \leq Wに対して成り立ちます。
証明
kについての帰納法で示す。
k=0のとき、a^{(0)}_{i,j} = 0 \ (\forall i, j) であるので、b^{(0)}_{i, j} = c^{(0)}_{i, j} = e^{(0)}_{i, j} \ (\forall i, j)となる。k=nでb^{(k)}_{i, j} - b^{(k)}_{i, j-1} = e^{(k)}_{i, j}かつc^{(k)}_{i, j} - c^{(k)}_{i-1, j} = e^{(k)}_{i, j}が成り立つとする。
(4)より、
b^{(n+1)}_{i, j} - b^{(n+1)}_{i, j-1} =
\left\{
\begin{array}{ll}
b^{(n)}_{i, j} - b^{(n)}_{i, j-1}+x_{n+1} & ((i, j) = (u_{n+1}, l_{n+1}), \ (d_{n+1}+1, r_{n+1}+1)) \\
b^{(n)}_{i, j} - b^{(n)}_{i, j-1} - x_{n+1} & ((i, j) = (u_{n+1}, r_{n+1}+1), \ (d_{n+1}+1, l_{n+1})) \\
b^{(n)}_{i, j} - b^{(n)}_{i, j-1} & (\text{otherwise})
\end{array}
\right.
となるので、帰納法の仮定よりb^{(n+1)}_{i, j} - b^{(n)}_{i, j} = e^{(n+1)}_{i, j} \ (\forall i, j)が成り立つ。また(5)より、
c^{(n+1)}_{i, j} - c^{(n+1)}_{i-1, j} =
\left\{
\begin{array}{ll}
c^{(n)}_{i, j} - c^{(n)}_{i, j-1}+x_{n+1} & ((i, j) = (u_{n+1}, l_{n+1}), \ (d_{n+1}+1, r_{n+1}+1)) \\
c^{(n)}_{i, j} - c^{(n)}_{i, j-1} - x_{n+1} & ((i, j) = (u_{n+1}, r_{n+1}+1), \ (d_{n+1}+1, l_{n+1})) \\
c^{(n)}_{i, j} - c^{(n)}_{i, j-1} & (\text{otherwise})
\end{array}
\right.
となるので、帰納法の仮定よりc^{(n+1)}_{i, j} - c^{(n)}_{i, j} = e^{(n+1)}_{i, j} \ (\forall i, j)が成り立つ。
従って、すべての1\leq k \leq Q, \ 1\leq i \leq H, \ 1\leq j \leq Wに対して
\begin{align}
b^{(k)}_{i, j} = \sum_{l=1}^{j}\left(b^{(k)}_{i, l} - b^{(k)}_{i, l-1}\right) = \sum_{l=1}^{j}e^{(k)}_{i, l}, \\
c^{(k)}_{i, j} = \sum_{l=1}^{j}\left(c^{(k)}_{l, j} - c^{(l)}_{l-1, j}\right) = \sum_{l=1}^{j}e^{(k)}_{l, j}
\end{align}
が成り立ちます。
各操作に対してe^{(k-1)}_{i, j}からe^{(k)}_{i, j}を求める部分の計算量がO(1)なので
e^{(Q)}_{i, j}を求めるのにO(Q)の計算量がかかり、またe^{(Q)}_{i, j}からb^{(Q)}_{i, j}, c^{(Q)}_{i, j} を求める部分にO(HW)の計算量がかかるので、b^{(Q)}_{i, j}, c^{(Q)}_{i, j} を求めるために必要な計算量はO(Q+HW)となります。
また、
a^{(Q)}_{i, j} = a^{(Q)}_{i, j-1} + c^{(Q)}_{i, j}, \ a^{(Q)}_{i, j} = a^{(Q)}_{i-1, j} +b^{(Q)}_{i, j}
であるので、b^{(Q)}_{i, j}, c^{(Q)}_{i, j}を用いてa^{(Q)}_{i, j}を求めるために必要な計算量はO(HW)となります。
以上よりO(Q+HW)でQ回の操作後のa^{(Q)}_{i, j}を求めることができました。
なお、各iについてb^{(Q)}_{i, 1}, ..., b^{(Q)}_{i, W}を求めるごとにe^{(Q)}_{i, j}をb^{(Q)}_{i, j}に順次更新し、その上で(7)の計算を行えば、得られるc^{(Q)}_{i, j}は求めたいa^{(Q)}_{i, j}と一致します。b^{(Q)}_{i, j}, c^{(Q)}_{i, j}を別々に求めるよりも、このような実装の方が楽かと思います。
実装
Ruby と Pythonでそれぞれclass を用いてimos法を実装したものです。上記での説明とは違い、0-index での実装となっていることには注意が必要です。また、値を更新する区間は半開区間としていることも注意が必要です。
pythonで実装したコード (1次元、2次元両方)
Rubyで実装したコード (1次元のみ)
応用問題
https://atcoder.jp/contests/abc338/tasks/abc338_d
この問題をコンテスト中に解けなかったことがきっかけとなって、この記事を書きました。
参考文献
より詳しい解説はこちら。特殊な座標系へのいもす法の適用なども書かれています。
https://imoz.jp/algorithms/imos_method.html
２次元のいもす法を可視化されています。
https://zenn.dev/reputeless/articles/article-2d-imos
こちらの方も数式を用いて説明されています。
https://zenn.dev/iharuoru/articles/c464535ced4b3b
Discussion