いもす法をpythonのclassで書いてみた

はじめに

いもす法は割と競プロの中ではメジャーなアルゴリズムですが、毎回書くよりclassで書いたほうが楽じゃないのと思ったけど誰もpythonで書いてなさそうだったので自分で書いてみてみたってやつです。

いもす法について

いもす法について説明は省きますが以下のページが詳しいです。

本題

コード

以下がコードです。

class Imos():
    """
    0次1次元いもす法

    Parameters
    ----------
    n : int
        配列の大きさ
    default : int
        リストの初期値
    """

    def __init__(self, n: int, default: int=0):
        self.n = n
        self.data = [default] + [0] * n

    def add(self, start: int, end: int, value: int=1):
        # 区間[start, end]にvalueを加算する
        self.data[start] += value
        self.data[end + 1] -= value

    def get(self):
        # 配列を取得する
        res = [0] * self.n
        for i in range(self.n):
            res[i] = res[i - 1] + self.data[i]
        return res

用例

こんな感じで使えます。例題はABC035-C オセロです。

def resolve():
    n, q = map(int, input().split())
    imos = Imos(n)
    for _ in range(q):
        l, r = map(int, input().split())
        imos.add(l - 1, r - 1, 1)
    print("".join([str(x % 2) for x in imos.get()]))

# class Imos():
#     省略

if __name__ == "__main__":
    resolve()

2次元バージョン

典型90問の某問題に二次元いもす法が出てきたので、それ用に拡張したやつも置いておきます。

コード

class Imos2d():
    """
    0次2次元いもす法

    Parameters
    ----------
    h : int
        配列の高さ(yの幅)
    w : int
        配列の幅(xの幅)
    default : int
        リストの初期値
    """

    def __init__(self, h: int, w: int, default: int=0):
        self.h = h
        self.w = w
        self.data = [[0] * (w + 1) for _ in range(h + 1)]
        self.data[0][0] = default

    def add(self, start_y: int, start_x: int, end_y: int, end_x: int, value: int=1):
        # 区間[(start_x, start_y), (end_x, end_y)]にvalueを加算する
        self.data[start_y][start_x] += value
        self.data[start_y][end_x + 1] -= value
        self.data[end_y + 1][start_x] -= value
        self.data[end_y + 1][end_x + 1] += value

    def get(self):
        # 配列を取得する
        res = [[0] * self.w for _ in range(self.h)]
        for h in range(self.h):
            for w in range(self.w):
                res[h][w] = res[h][w - 1] + self.data[h][w]
        for w in range(self.w):
            for h in range(1, self.h):
                res[h][w] += res[h - 1][w]
        return res

2次元バージョンの用例

例題はyukicoder No.60 魔法少女です。

def resolve():
    import sys
    input = sys.stdin.readline
    n, k = map(int, input().split())
    enemies = [list(map(int, input().split())) for _ in range(n)]
    imos = Imos2d(1501, 1501)
    for _ in range(k):
        x, y, w, h, d = map(int, input().split())
        imos.add(y + 500, x + 500, y + h + 500, x + w + 500, d)
    damage = imos.get()
    ans = 0
    for x, y, hp in enemies:
        ans += max(hp - damage[y + 500][x + 500], 0)
    print(ans)


# class Imos2d():
#     省略


if __name__ == '__main__':
    resolve()

典型90問のものも終了してしばらくしたら公開する予定です。

おわりに

毎回自分で書いてごちゃごちゃになりがちなものはclassで書いてライブラリ化しておくのも良いなと思った所感です。

徒然

思いついたら書きます