クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「3-3 BITの実装」(p.161)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「3-3 BITの実装」(p.161)
のコードをPythonで書き直したものとなっています。

蟻本のコード

# [1, n]
n = int(input())
bit = [0] * (n + 1)

def sum_(i):
    s = 0
    while i > 0:
        s += bit[i]
        i -= i & -i
    return s

def add(i, x):
    while i <= n:
        bit[i] += x
        i += i & -i

classを使ったコード

class BinaryIndexedTree:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)
        
    def sum(self, i):
        s = 0
        while i > 0:
            s += self.bit[i]
            i -= i & -i
        return s

    def add(self, i, x):
        while i <= self.n:
            self.bit[i] += x
            i += i & -i