ベスパリブ

プログラミングを主とした日記・備忘録です。ベスパ持ってないです。

Pythonでいもす法

いもす法とは

いもす法(imos法)とは、長さNの1次元配列において、「ある連続する区間に、ある数vを足す」という操作をK回繰り返した結果を、計算量O(N+K)で高速に構築できる方法。

たとえば、「ある連続する区間に、1を足す」という操作を4回したい次の図のようなものを考えた場合、いもす法を使えば最終結果を高速に求めることができます。

f:id:takeg:20190828105337j:plain
いもす法の例

2次元や3次元に拡張できるらしいです(まだ使ったことない)。2次元に拡張バージョンは本家の説明を読むとなるほどなぁとなれます。

いもす法の手順

いもす法の手順は以下の2ステップです。

  • ステップ1: 加算処理
    • 区間 [l, r] に v を加算したいとき、
      • l 番目の値に v を加算する
      • r+1 番目の値に -v を加算する
  • ステップ2: 累積和
    • 最後に累積和をすると、最終結果を得られる

例1

長さ7のリストにおいて、

という3つの操作をしたときの最終結果を得たいとします。

ステップ1: 加算処理

  • 初期状態
i 0 1 2 3 4 5 6
要素の値 0 0 0 0 0 0 0
N = 7
data = [0] * N

区間[1, 3]に2を加算

i 0 1 2 3 4 5 6
要素の値 0 2 0 0 -2 0 0

l = 1, r = 3 なので、

  • data[l] に 2 を加算
  • data[r+1] に -2 を加算

をします。

l, r = 1, 3
v = 2
data[l] += v
if r+1 != N:
    data[r+1] +=-v

区間[3, 3]に3を加算

i 0 1 2 3 4 5 6
要素の値 0 2 0 3 -5 0 0
l, r = 3, 3
v = 3
data[l] += v
if r+1 != N:
    data[r+1] +=-v

区間[0, 5]に1を加算

i 0 1 2 3 4 5 6
要素の値 1 2 0 3 -5 0 -1
l, r = 0, 5
v = 1
data[l] += v
if r+1 != N:
    data[r+1] +=-v

ステップ2: 累積和

最後に、累積和をすると最終結果を得ることができます。

i 0 1 2 3 4 5 6
要素の値 1 2 0 3 -5 0 -1
累積和 1 3 3 6 1 1 0
ruiseki = [0] * N  # 累積和用リスト
ruiseki[0] = data[0]
for i in range(1, N):
    ruiseki[i] = ruiseki[i-1] + data[i]

# 最終結果を出力
print(ruiseki)
# [1, 3, 3, 6, 1, 1, 0]

例2: AtCoder Beginner Contest 035: C - オセロ

黒の面に0、白の面に1が書かれた N 個のオセロの駒が、どの駒も黒の面が上を向くように一列に並べられています。その後、ある区間にある駒を全て裏返すという操作が Q 回だけ行なわれました。 具体的には i 回目の操作においては、左から l_i 番目の駒から r_i 番目の駒までの駒全てが裏返されました。

最終的な盤面を求めてください。

# -*- coding:utf-8 -*-

def solve():
    # 入力処理
    N, Q = list(map(int, input().split()))

    # オセロの駒の初期化(最初はすべて黒)
    s = [0] * (N+1)
 
    # いもす法の加算処理
    for _ in range(Q):
        l, r = list(map(int, input().split()))  # 区間[l, r]の入力を受け取る
        s[l] += 1
        if r+1 != N+1:
            s[r+1] -= 1
    
    # いもす法の累積和(構築処理)
    ruiseki = [0] * (N+1)
    for i in range(1, N+1):
        ruiseki[i] += ruiseki[i-1] + s[i]
        if ruiseki[i]%2 == 0:
            # 裏返した回数が偶数なら、黒のまま
            ruiseki[i] = 0
        else:
            # 裏返した回数が奇数なら、白
            ruiseki[i] = 1
    
    # 答えを出力
    ans = "".join(list(map(str, ruiseki[1:])))
    print(ans)
 
 
if __name__ == "__main__":
    solve()

感想

r+1 に -v を加算することで、累積和したときによしなに調整されるというアイデアアルゴリズムでした。シンプルだけどすごい。

いもす法で解く問題

編集履歴

  • 2020/04/24: 「ステップ1: 加算処理」と「ステップ2: 累積和」を微修正。