Python Imos法ライブラリ

Imos法ライブラリ

imos法のPythonライブラリです。
imos法はあるの区間に対して処理を行うクエリがたくさん与えられるときに使うアルゴリズムです。詳しくはいもすさんによる解説記事を見てください。

from itertools import accumulate
class Imos:
    def __init__(self, n):
        self.A = [0]*(n+1)
        self.n = n

    def __call__(self):
        dist = list(accumulate(self.A))[:self.n]
        self.A = [0] * (self.n+1)
        return dist

    def query(self, l, r):
        self.A[l] += 1
        self.A[r] -= 1

解説

... リストのサイズを指定します。
... 区間 に対するクエリを入力に取ります。
... クエリの反映を行いリスト状態を返します。

使用例

使用したコードがあれば適宜アップデートしていきます。

AtCoder ABC 035 C - オセロ

from itertools import accumulate
class Imos:
    # Folding

n, q, *LR = map(int, open(0).read().split())
imos = Imos(n)
for l, r in zip(LR[::2], LR[1::2]):
    imos.query(l-1, r)
print(*(x%2 for x in imos()), sep='')

AtCoder ABC 014 C - AtColor

from itertools import accumulate
class Imos:
    # Folding

n, *AB = map(int, open(0).read().split())
imos = Imos(1_000_001)
for l, r in zip(AB[::2], AB[1::2]):
    imos.query(l, r+1)
print(max(imos()))

AtCoder 東京海上日動 プログラミングコンテスト2020 C - Lamps

# Python3 is TLE! Execute by Pypy3
from itertools import accumulate
class Imos:
    # Folding

n, k, *A = map(int, open(0).read().split())
imos = Imos(n)
for _ in range(min(100, k)):
    for i in range(n):
        imos.query(max(0, i-A[i]), min(i+A[i]+1, n))
    A = imos()
print(*A)