Binary Indexed Treeの区間合計の定数倍改善手法の紹介

## はじめに

はじめましての人ははじめまして。[alumi](https://x.com/alumi_02)と申します。
少し前に書いた「[Binary Indexed Treeの一点参照の定数倍改善手法の紹介](https://qiita.com/alumite14/items/4a1b6143f2049eb8e100)」という記事で、一点取得の高速化に関する考察をしましたが、区間合計でも同様に定数倍高速化ができ、しかも前述の一点取得は区間合計手法の特殊な場合にまとめられることに気づいたので、解説と実装結果を紹介します。

## 結論

- この実装をすると、`prefix_sum`の差による区間合計の実装に比べて定数倍改善できます。
- 一点取得を`range_sum(i, i + 1)`で実装しても、前回紹介した手法と同程度に高速です。
- ほぼ同じ考え方で区間加算・区間合計BITについても区間合計・一点取得を定数倍改善できます。

## 実装と手法の解説

早速実装を見ましょう。
注意点として、一般的な実装と同様に逆元が定義されている必要があります（参考：[BIT の定数倍高速化について - ecasdqina-cp](https://scrapbox.io/ecasdqina-cp/BIT_%E3%81%AE%E5%AE%9A%E6%95%B0%E5%80%8D%E9%AB%98%E9%80%9F%E5%8C%96%E3%81%AB%E3%81%A4%E3%81%84%E3%81%A6)）

```python
    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res = 0
        while l < r:
            res += self.tree[r]
            r -= r & -r
        while r < l:
            res -= self.tree[l]
            l -= l & -l
        return res
```

一般的な`prefix_sum`の差による区間合計では、index0に近い側のblockが共通となり、同じ計算を繰り返す場合があります（下図）。ただ図を見ればわかる通り、ある区間$[l,r)$の区間合計に必要なblockは、**その区間をすべて被覆するようなblockのうちlsbが最も小さいものの下にあるblockだけ**です。それより上にあるblockは計算に不要です（下図$[9,14)$ならtree index 12, 14が相当。具体的な証明は上リンクを参照ください）。

![range_sum01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665424/5ccaf69e-1d42-3092-810c-3f4b7b4c134b.png)

そこで$l$を直上で被覆するようなblockとなるまで、$r$からlsbを引きまくりその間の区間を足し、引いた時点での$r$を直上で被覆するようなblockとなるまで、$l$からlsbを引きまくりその間の区間を引くと、ほしい区間合計が得られます。

### 一点取得高速化手法との関係 

前回の記事で紹介した一点取得高速化手法の実装は以下のようになります。
```python
    def __getitem__(self, i: int):
        j = i + 1
        k = 1
        res = self.tree[j]
        while j & k == 0:
            res -= self.tree[j - k]
            k <<= 1
        return res
```

今回の手法において$l=i, r=i+1$とすると、まず最初のwhileは一度しか通らないことがわかります。なぜならlsbは必ず1以上なので、一回lsbを引かれると即座に`l < r`を満たさなくなるためです。
よって最初のwhileは`res = self.tree[i + 1]`で置き換えられます。

2つ目のwhileでは$l$が$r$と同じlsbを持つまでlsbを引いています（図を見ると$r$と同じ高さを目指しているのがわかりやすいです）。これは一点取得手法における`while j & k == 0`, `k <<= 1` の実装と同値です（$j = i+1 = r$として`while r & k == 0`と見れば$r$のlsbが見つかるまで走査していることがわかる）。

以上から前回の記事で紹介した一点取得高速化手法は、今回の区間加算における特殊な場合として見ることができます。よって区間加算だけ今回の手法で実装してしまえば、一点取得については一般的な`range_sum(i, i + 1)`のままでも前回手法と同程度に高速化できます。うれしい。

<details><summary>実装</summary>

```python
class BinaryIndexedTree:
    __slots__ = ("n", "tree")

    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [0] * (self.n + 1)

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[j] += ins.tree[i]
        return ins

    def prefix_sum(self, i: int):
        """Prefix sum of [0, i) (0-indexed)"""
        res = 0
        while i > 0:
            res += self.tree[i]
            i -= i & -i
        return res

    def add(self, i: int, x) -> None:
        """Add x at i (0-indexed)"""
        i += 1
        while i <= self.n:
            self.tree[i] += x
            i += i & -i

    def update(self, i: int, x):
        """Update i to x (0-indexed)"""
        self.add(i, x - self[i])

    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res = 0
        while l < r:
            res += self.tree[r]
            r -= r & -r
        while r < l:
            res -= self.tree[l]
            l -= l & -l
        return res

    def __getitem__(self, i: int):
        return self.range_sum(i, i + 1)

    def __setitem__(self, i: int, val):
        self.update(i, val)

    def __iter__(self):
        yield from (self[i] for i in range(self.n))

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({list(self)})"

```

</details>

## 区間加算・区間合計BITへの適用

前回同様に区間加算・区間合計BITへ適用できないか考えてみます。
一般的な実装では、`range_sum`は次のような実装がされています。

```python
    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        return self.prefix_sum(r) - self.prefix_sum(l)
```

また区間加算・区間合計BITにおける`prefix_sum`は次のような実装がされています。

```python
    def prefix_sum(self, i: int):
        """Prefix sum of [0, i) (0-indexed)"""
        return self._sum_partial(0, i) + self._sum_partial(1, i) * i
```

ここで`_sum_partial`は一点取得BITにおけるprefix_sumとほぼ同様の実装で、第一引数に応じて計算で用いるtreeを切り替えている以外に違いはありません。

```python
    def _sum_partial(self, p: int, i: int):
        res = 0
        while i > 0:
            res += self.tree[p][i]
            i -= i & -i
        return res
```

よって式変形で上手いこと`_sum_partial(p, r) - _sum_partial(p, l)`の形をつくってやれば、この部分を高速化した`renge_sum`に置き換えることで高速化できそうです。

区間加算・区間合計BITの`range_sum`は次のように式変形できます。

```math
\begin{eqnarray*}
\text{range_sum}(l,\ r) & = & \text{prefix_sum}(r)-\text{prefix_sum}(l) \\\\
& = & \text{sum_partial}(0,\ r) + \text{sum_partial}(1,\ r) * r \\
& - & \text{sum_partial}(0,\ l) - \text{sum_partial}(1,\ l) * l \\\\
& = & \text{sum_partial}(0,\ r) - \text{sum_partial}(0,\ l) \\
& + & \text{sum_partial}(1,\ r) * r - \text{sum_partial}(1,\ l) * l \\\\
& = & \text{sum_partial}(0,\ r) - \text{sum_partial}(0,\ l) \\
& + & \text{sum_partial}(1,\ r) * l - \text{sum_partial}(1,\ l) * l \\
& + & \text{sum_partial}(1,\ r) * (r - l) \\\\
& = & \underline{\text{sum_partial}(0,\ r) - \text{sum_partial}(0,\ l)} \\
& + & \big\lbrace\underline{\text{sum_partial}(1,\ r) - \text{sum_partial}(1,\ l)}\big\rbrace * l \\
& + & \text{sum_partial}(1,\ r) * (r - l)
\end{eqnarray*}
```

下線部でほしい形が作れていますね。ちなみに$\text{sum\_partial}(1,\ r) - \text{sum\_partial}(1,\ l)$の形を作るために$\text{sum\_partial}(1,\ r) * r = \text{sum\_partial}(1,\ r) * l + \text{sum\_partial}(1,\ r) * (r-l)$と分離しています。
この部分を一点加算・区間合計BITと同様のrange_sum（下の`_range_sum_partial`）で置き換えてやればいいです。

```python
    def _range_sum_partial(self, p: int, l: int, r: int):
        res = 0
        while l < r:
            res += self.tree[p][r]
            r -= r & -r
        while r < l:
            res -= self.tree[p][l]
            l -= l & -l
        return res

    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res0 = self._range_sum_partial(0, l, r)
        res1 = self._range_sum_partial(1, l, r)
        return res0 + res1 * l + self._sum_partial(1, r) * (r - l)
```

このままでもいいのですが、`res0`, `res1`は`_range_sum_partial`で同じ区間$[l,r)$について計算しているのでまとめられそうですね。
更に`_range_sum_partial`はどうせここでしか使わないので`range_sum`内にぶち込んでしまいましょう。注意点としては最後にかける`l, r`と`res0`, `res1`の計算で使う`l, r`を区別する必要があるので、`tmp_l`, `tmp_r`を用意しています。

```python
    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res0 = res1 = 0
        tmp_l, tmp_r = l, r
        while tmp_l < tmp_r:
            res0 += self.tree[0][tmp_r]
            res1 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        while tmp_r < tmp_l:
            res0 -= self.tree[0][tmp_l]
            res1 -= self.tree[1][tmp_l]
            tmp_l -= tmp_l & -tmp_l
        return res0 + res1 * l + self._sum_partial(1, r) * (r - l)
```

ちなみに前回記事の区間加算・区間合計BITの一点取得高速化の提案は次のような感じでした。

```python
    def __getitem__(self, i: int):
        j = i + 1
        k = 1
        res0 = self.tree[0][j]
        res1 = self.tree[1][j]
        while j & k == 0:
            res0 -= self.tree[0][j - k]
            res1 -= self.tree[1][j - k]
            k <<= 1
        return res0 + res1 * i + self._sum_partial(1, j)
```

よく似ていますね。実際$l,\ r = i,\ i+1$とすれば[一点取得高速化手法との関係](#一点取得高速化手法との関係)で述べた通り同値であることが示せます。区間加算・区間合計BITにおいても前回の記事の内容は今回の特殊系であることがわかりますね。

これで完了でもいいのですが、実は`_sum_partial(1, r)`の部分がまだまとめられます。
`_sum_partial(1, r)`では、**$r$が$0$になるまでlsbを引きながらresに繰り返す足す**計算をします。
`range_sum`の`res1`を見ると、**$r$が$l$を下回るまでlsbを引きながら`res1`に繰り返す足す**計算をしています。つまり`_sum_partial(1, r)`の計算を途中まで行ってくれているということですね。

```python
    def _sum_partial(self, p: int, i: int):
        res = 0
        # こっちはr <= 0まで
        while i > 0:
            res += self.tree[p][i]
            i -= i & -i
        return res

    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res0 = res1 = 0
        tmp_l, tmp_r = l, r
        # こっちはr <= lまで
        while tmp_l < tmp_r:
            res0 += self.tree[0][tmp_r]
            res1 += self.tree[1][tmp_r] 
            tmp_r -= tmp_r & -tmp_r
		...
```

なので`range_sum`において1つ目のwhileから出た時点の`res1`を`res2`として取っておき、最後に残った$r$を$0$になるまでlsbを引きながら`res2`に足していけば、`res2`は`_sum_partial(1, r)`と同じ値になります。
以上をまとめると最終形はこんな感じになります。あんまりきれいではないですね。

```python
    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res0 = res1 = 0
        tmp_l, tmp_r = l, r
        while tmp_l < tmp_r:
            res0 += self.tree[0][tmp_r]
            res1 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        res2 = res1
        while tmp_r < tmp_l:
            res0 -= self.tree[0][tmp_l]
            res1 -= self.tree[1][tmp_l]
            tmp_l -= tmp_l & -tmp_l
        while tmp_r > 0:
            res2 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        return res0 + res1 * l + res2 * (r - l)
```

ちなみに`prefix_sum`も同様のまとめ方ができます。つまり
```python
    def prefix_sum(self, i: int):
        return self._sum_partial(0, i) + self._sum_partial(1, i) * i
```

としているものを
```python
    def prefix_sum(self, i: int):
        res0 = res1 = 0
        tmp_i = i
        while tmp_i > 0:
            res0 += self.tree[0][tmp_i]
            res1 += self.tree[1][tmp_i]
            tmp_i -= tmp_i & -tmp_i
        return res0 + res1 * i
```

とまとめることができます。こうすると`_sum_partial`は他のメソッドで使われないのでまるっと消せます[^1]。一応これについても速度チェックしておきましょうか。

[^1]: この形式の`prefix_sum`は今回の手法の`range_sum`の特殊系となるので、`range_sum(0, i)`と実装してもいい

## 速度チェック

### 一点加算・区間合計BIT｜区間合計

ランダムな配列Aにより初期化したBITインスタンスを生成し、ランダムな区間合計クエリに対してかかる時間を比較します。

<details><summary>テストコード</summary>

```python
class BinaryIndexedTree:
    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [0] * (self.n + 1)

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[j] += ins.tree[i]
        return ins

    def prefix_sum(self, i: int):
        """Prefix sum of [0, i) (0-indexed)"""
        res = 0
        while i > 0:
            res += self.tree[i]
            i -= i & -i
        return res

    def range_sum_prefix_diff(self, l: int, r: int):
        return self.prefix_sum(r) - self.prefix_sum(l)

    def range_sum_two_while(self, l: int, r: int):
        res = 0
        while l < r:
            res += self.tree[r]
            r -= r & -r
        while r < l:
            res -= self.tree[l]
            l -= l & -l
        return res

import random
import timeit
import matplotlib.pyplot as plt

# テストパラメータ
N = 10**5  # 配列の長さ
Q = 10**5  # クエリの数
random.seed(42)  # Answer to the Ultimate Question of Life, the Universe, and Everything

# ランダムな要素の配列Aを作成し、BITを初期化
A = [random.randint(-(10**9), 10**9) for _ in range(N)]
bit = BinaryIndexedTree.from_array(A)

# ランダムなクエリを生成し、test_caseに格納
test_case = [(random.randrange(N), random.randrange(N)) for _ in range(Q)]
test_case = [(min(l, r), max(l, r)) for l, r in test_case]


# それぞれのメソッドをtest_caseで速度チェック
def test_range_sum_prefix_diff():
    for l, r in test_case:
        bit.range_sum_prefix_diff(l, r)


def test_range_sum_two_while():
    for l, r in test_case:
        bit.range_sum_two_while(l, r)

# loop回の平均を取る
loop = 100
time_prefix_diff = timeit.timeit(test_range_sum_prefix_diff, number=loop)
time_two_while = timeit.timeit(test_range_sum_two_while, number=loop)

# 結果を棒グラフで表示
methods = ["prefix_diff", "two_while"]
times = [time_prefix_diff / loop, time_two_while / loop]
y_positions = [0, 0.6]  # Y軸の位置を明示的に指定
plt.barh(y_positions, times, color=["blue", "green"], height=0.4)
plt.yticks(y_positions, methods)  # Y軸のラベルと位置を設定
plt.xlabel("Execution Time (sec)")
plt.title("Range Sum Comparison")
plt.show()

```

</details>

#### 結果

![result_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665424/df78620e-9d24-e9c4-3b7c-4c4117d6eba5.png)

わずかですが`prefix_diff`（一般的な実装）より`two_while`（今回の実装）が速い結果が出ています。

### 一点加算・区間合計BIT｜一点取得

ランダムな一点取得に対してかかる時間を比較します。

<details><summary>テストコード</summary>

```python
class BinaryIndexedTree:
    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [0] * (self.n + 1)

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[j] += ins.tree[i]
        return ins

    def prefix_sum(self, i: int):
        res = 0
        while i > 0:
            res += self.tree[i]
            i -= i & -i
        return res

    def range_sum(self, l: int, r: int):
        res = 0
        while l < r:
            res += self.tree[r]
            r -= r & -r
        while r < l:
            res -= self.tree[l]
            l -= l & -l
        return res

    def getitem_two_while(self, i: int):
        return self.range_sum(i, i + 1)

    def getitem_prefix_diff(self, i: int):
        return self.prefix_sum(i + 1) - self.prefix_sum(i)

    def getitem_one_while(self, i: int):
        j = i + 1
        k = 1
        res = self.tree[j]
        while j & k == 0:
            res -= self.tree[j - k]
            k <<= 1
        return res


import random
import timeit
import matplotlib.pyplot as plt

# テストパラメータ
N = 10**5  # 配列の長さ
Q = 10**5  # クエリの数
random.seed(42)  # Answer to the Ultimate Question of Life, the Universe, and Everything

# ランダムな要素の配列Aを作成し、BITを初期化
A = [random.randint(-(10**9), 10**9) for _ in range(N)]
bit = BinaryIndexedTree.from_array(A)

# ランダムなクエリを生成し、test_caseに格納
test_case = [random.randrange(N) for _ in range(Q)]


# それぞれのメソッドをtest_caseで速度チェック
def test_getitem_prefix_diff():
    for i in test_case:
        bit.getitem_prefix_diff(i)


def test_getitem_two_while():
    for i in test_case:
        bit.getitem_two_while(i)


def test_getitem_one_while():
    for i in test_case:
        bit.getitem_one_while(i)

# loop回の平均を取る
loop = 100
time_prefix_diff = timeit.timeit(test_getitem_prefix_diff, number=loop)
time_two_while = timeit.timeit(test_getitem_two_while, number=loop)
time_one_while = timeit.timeit(test_getitem_one_while, number=loop)


# 結果を棒グラフで表示
methods = ["prefix_diff", "two_while", "one_while"]
times = [time_prefix_diff / loop, time_two_while / loop, time_one_while / loop]
y_positions = [0, 0.6, 1.2]  # Y軸の位置を明示的に指定
plt.barh(y_positions, times, color=["blue", "green", "red"], height=0.4)
plt.yticks(y_positions, methods)  # Y軸のラベルと位置を設定
plt.xlabel("Execution Time (sec)")
plt.title("Get Item Comparison")
plt.show()
```

</details>


#### 結果

![result_2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665424/a809d15b-72a1-651d-4946-8be2c60c64b2.png)

`prefix_diff`（一般的な実装）よりもそれ以外が相当早くなっていますね。`one_while`（前回の実装）のほうが`two_while`（今回の実装）よりもわずかに早いです。実装を簡潔にする以上に高速化のニーズが高いならどっちも実装すると良いでしょうけど、この程度の差であれば`two_while`で実装をまとめたほうが嬉しみがありそうですね。

### 区間加算・区間合計BIT｜区間合計

ランダムな配列Aにより初期化したBITインスタンスを生成し、ランダムな区間加算クエリを与えてから、ランダムな区間合計クエリに対してかかる時間を比較します。

<details><summary>テストコード</summary>

```python
class BinaryIndexedTree:
    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [[0] * (self.n + 1), [0] * (self.n + 1)]

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[0][1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[0][j] += ins.tree[0][i]
        return ins

    def _sum_partial(self, p: int, i: int):
        res = 0
        while i > 0:
            res += self.tree[p][i]
            i -= i & -i
        return res

    def prefix_sum(self, i: int):
        """Prefix sum of [0, i) (0-indexed)"""
        return self._sum_partial(0, i) + self._sum_partial(1, i) * i

    def _add_partial(self, p: int, i: int, x) -> None:
        while i <= self.n:
            self.tree[p][i] += x
            i += i & -i

    def add_range(self, l: int, r: int, x) -> None:
        """Add x at [l, r) (0-indexed)"""
        l += 1
        self._add_partial(0, l, -x * (l - 1))
        self._add_partial(0, r, x * r)
        self._add_partial(1, l, x)
        self._add_partial(1, r, -x)

    def range_sum_prefix_diff(self, l: int, r: int):
        return self.prefix_sum(r) - self.prefix_sum(l)

    def _range_sum_partial(self, p: int, l: int, r: int):
        res = 0
        while l < r:
            res += self.tree[p][r]
            r -= r & -r
        while r < l:
            res -= self.tree[p][l]
            l -= l & -l
        return res

    def range_sum_use_rsp(self, l: int, r: int):
        res0 = self._range_sum_partial(0, l, r)
        res1 = self._range_sum_partial(1, l, r)
        return res0 + res1 * l + self._sum_partial(1, r) * (r - l)

    def range_sum_three_while(self, l: int, r: int):
        res0 = res1 = 0
        tmp_l, tmp_r = l, r
        while tmp_l < tmp_r:
            res0 += self.tree[0][tmp_r]
            res1 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        res2 = res1
        while tmp_r < tmp_l:
            res0 -= self.tree[0][tmp_l]
            res1 -= self.tree[1][tmp_l]
            tmp_l -= tmp_l & -tmp_l
        while tmp_r > 0:
            res2 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        return res0 + res1 * l + res2 * (r - l)


import random
import timeit
import matplotlib.pyplot as plt

# テストパラメータ
N = 10**5  # 配列の長さ
Q = 10**5  # クエリの数
random.seed(42)  # Answer to the Ultimate Question of Life, the Universe, and Everything

# ランダムな要素の配列Aを作成し、BITを初期化
A = [random.randint(-(10**9), 10**9) for _ in range(N)]
bit = BinaryIndexedTree.from_array(A)
# tree1側にも値を格納
for _ in range(1000):
    l = random.randrange(N + 1)
    r = random.randrange(N + 1)
    if l == r:
        continue
    l, r = min(l, r), max(l, r)
    x = random.randint(-10**9, 10**9)
    bit.add_range(l, r, x)


# ランダムなクエリを生成し、test_caseに格納
test_case = [(random.randrange(N), random.randrange(N)) for _ in range(Q)]
test_case = [(min(l, r), max(l, r)) for l, r in test_case]


# それぞれのメソッドをtest_caseで速度チェック
def test_range_sum_prefix_diff():
    for l, r in test_case:
        bit.range_sum_prefix_diff(l, r)


def test_range_sum_use_rsp():
    for l, r in test_case:
        bit.range_sum_use_rsp(l, r)


def test_range_sum_three_while():
    for l, r in test_case:
        bit.range_sum_three_while(l, r)


# loop回の平均を取る
loop = 100
time_prefix_diff = timeit.timeit(test_range_sum_prefix_diff, number=loop)
time_use_rsp = timeit.timeit(test_range_sum_use_rsp, number=loop)
time_three_while = timeit.timeit(test_range_sum_three_while, number=loop)


# 結果を棒グラフで表示
methods = ["prefix_diff", "use_rsp", "three_while"]
times = [time_prefix_diff / loop, time_use_rsp / loop, time_three_while / loop]
y_positions = [0, 0.6, 1.2]  # Y軸の位置を明示的に指定
plt.barh(y_positions, times, color=["red", "blue", "green"], height=0.4)
plt.yticks(y_positions, methods)  # Y軸のラベルと位置を設定
plt.xlabel("Execution Time (sec)")
plt.title("Range Sum Comparison")
plt.show()
```

</details>

#### 結果

![result_3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665424/31ccc9fa-15fb-b2d6-9246-5e0a1c11b88a.png)

ちゃんと`three_while`（今回の実装）が最速になってますね。ただ意外にも`use_rsp`（`_range_sum_partial`を用意して実装した実装）が`prefix_diff`（一般的な実装）に劣る結果になってますね。へんなの。

### 区間加算・区間合計BIT｜一点取得

ランダムな一点取得に対してかかる時間を比較します。

<details><summary>テストコード</summary>

```python
class BinaryIndexedTree:
    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [[0] * (self.n + 1), [0] * (self.n + 1)]

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[0][1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[0][j] += ins.tree[0][i]
        return ins

    def _sum_partial(self, p: int, i: int):
        res = 0
        while i > 0:
            res += self.tree[p][i]
            i -= i & -i
        return res

    def prefix_sum(self, i: int):
        return self._sum_partial(0, i) + self._sum_partial(1, i) * i

    def _add_partial(self, p: int, i: int, x) -> None:
        while i <= self.n:
            self.tree[p][i] += x
            i += i & -i

    def add_range(self, l: int, r: int, x) -> None:
        """Add x at [l, r) (0-indexed)"""
        l += 1
        self._add_partial(0, l, -x * (l - 1))
        self._add_partial(0, r, x * r)
        self._add_partial(1, l, x)
        self._add_partial(1, r, -x)

    def range_sum_prefix_diff(self, l: int, r: int):
        return self.prefix_sum(r) - self.prefix_sum(l)

    def range_sum_three_while(self, l: int, r: int):
        res0 = res1 = 0
        tmp_l, tmp_r = l, r
        while tmp_l < tmp_r:
            res0 += self.tree[0][tmp_r]
            res1 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        res2 = res1
        while tmp_r < tmp_l:
            res0 -= self.tree[0][tmp_l]
            res1 -= self.tree[1][tmp_l]
            tmp_l -= tmp_l & -tmp_l
        while tmp_r > 0:
            res2 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        return res0 + res1 * l + res2 * (r - l)

    def getitem_one_while(self, i: int):
        j = i + 1
        k = 1
        res0 = self.tree[0][j]
        res1 = self.tree[1][j]
        while j & k == 0:
            res0 -= self.tree[0][j - k]
            res1 -= self.tree[1][j - k]
            k <<= 1
        return res0 + res1 * i + self._sum_partial(1, j)

    def getitem_three_while(self, i):
        return self.range_sum_three_while(i, i + 1)

    def getitem_prefix_diff(self, i):
        return self.range_sum_prefix_diff(i, i + 1)


import random
import timeit
import matplotlib.pyplot as plt

# テストパラメータ
N = 10**5  # 配列の長さ
Q = 10**5  # クエリの数
random.seed(42)  # Answer to the Ultimate Question of Life, the Universe, and Everything

# ランダムな要素の配列Aを作成し、BITを初期化
A = [random.randint(-(10**9), 10**9) for _ in range(N)]
bit = BinaryIndexedTree.from_array(A)
# tree1側にも値を格納
for _ in range(1000):
    l = random.randrange(N + 1)
    r = random.randrange(N + 1)
    if l == r:
        continue
    l, r = min(l, r), max(l, r)
    x = random.randint(-10**9, 10**9)
    bit.add_range(l, r, x)


# ランダムなクエリを生成し、test_caseに格納
test_case = [random.randrange(N) for _ in range(Q)]


# それぞれのメソッドをtest_caseで速度チェック
def test_getitem_prefix_diff():
    for i in test_case:
        bit.getitem_prefix_diff(i)


def test_getitem_two_while():
    for i in test_case:
        bit.getitem_three_while(i)


def test_getitem_one_while():
    for i in test_case:
        bit.getitem_one_while(i)

# loop回の平均を取る
loop = 100
time_prefix_diff = timeit.timeit(test_getitem_prefix_diff, number=loop)
time_three_while = timeit.timeit(test_getitem_two_while, number=loop)
time_one_while = timeit.timeit(test_getitem_one_while, number=loop)


# 結果を棒グラフで表示
methods = ["prefix_diff", "three_while", "one_while"]
times = [time_prefix_diff / loop, time_three_while / loop, time_one_while / loop]
y_positions = [0, 0.6, 1.2]  # Y軸の位置を明示的に指定
plt.barh(y_positions, times, color=["blue", "green", "red"], height=0.4)
plt.yticks(y_positions, methods)  # Y軸のラベルと位置を設定
plt.xlabel("Execution Time (sec)")
plt.title("Get Item Comparison")
plt.show()

```

</details>

#### 結果

![result_4.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665424/95298853-20d2-5eda-cb70-7c2bb7d0bb23.png)

ほぼ一点加算・区間合計BITのときの結果と同じですね。2倍以上早くなっています。強いて違いを言えば`one_while`と`three_while`との差がほぼなくなっている点でしょうか。この結果から`three_while`で実装をまとめても問題なさそうなことがわかりますね。

### 区間加算・区間合計BIT｜prefix_sum

ランダムな$i$に対する`prefix_sum`の計算にかかる時間を比較します。

<details><summary>テストコード</summary>

```python
class BinaryIndexedTree:
    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [[0] * (self.n + 1), [0] * (self.n + 1)]

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[0][1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[0][j] += ins.tree[0][i]
        return ins

    def _sum_partial(self, p: int, i: int):
        res = 0
        while i > 0:
            res += self.tree[p][i]
            i -= i & -i
        return res

    def prefix_sum_use_sp(self, i: int):
        return self._sum_partial(0, i) + self._sum_partial(1, i) * i

    def prefix_sum_one_while(self, i: int):
        res0 = res1 = 0
        tmp_i = i
        while tmp_i > 0:
            res0 += self.tree[0][tmp_i]
            res1 += self.tree[1][tmp_i]
            tmp_i -= tmp_i & -tmp_i
        return res0 + res1 * i

    def _add_partial(self, p: int, i: int, x) -> None:
        while i <= self.n:
            self.tree[p][i] += x
            i += i & -i

    def add_range(self, l: int, r: int, x) -> None:
        """Add x at [l, r) (0-indexed)"""
        l += 1
        self._add_partial(0, l, -x * (l - 1))
        self._add_partial(0, r, x * r)
        self._add_partial(1, l, x)
        self._add_partial(1, r, -x)

import random
import timeit
import matplotlib.pyplot as plt

# テストパラメータ
N = 10**5  # 配列の長さ
Q = 10**5  # クエリの数
random.seed(42)  # Answer to the Ultimate Question of Life, the Universe, and Everything

# ランダムな要素の配列Aを作成し、BITを初期化
A = [random.randint(-(10**9), 10**9) for _ in range(N)]
bit = BinaryIndexedTree.from_array(A)
# tree1側にも値を格納
for _ in range(1000):
    l = random.randrange(N + 1)
    r = random.randrange(N + 1)
    if l == r:
        continue
    l, r = min(l, r), max(l, r)
    x = random.randint(-10**9, 10**9)
    bit.add_range(l, r, x)


# ランダムなクエリを生成し、test_caseに格納
test_case = [random.randrange(N) for _ in range(Q)]


# それぞれのメソッドをtest_caseで速度チェック
def test_prefix_sum_use_sp():
    for i in test_case:
        bit.prefix_sum_use_sp(i)


def test_prefix_sum_one_while():
    for i in test_case:
        bit.prefix_sum_one_while(i)


# loop回の平均を取る
loop = 100
time_prefix_sum_use_sp = timeit.timeit(test_prefix_sum_use_sp, number=loop)
time_prefix_sum_one_while = timeit.timeit(test_prefix_sum_one_while, number=loop)


# 結果を棒グラフで表示
methods = ["use_sp", "one_while"]
times = [time_prefix_sum_use_sp / loop, time_prefix_sum_one_while / loop]
y_positions = [0, 0.6]  # Y軸の位置を明示的に指定
plt.barh(y_positions, times, color=["blue", "green"], height=0.4)
plt.yticks(y_positions, methods)  # Y軸のラベルと位置を設定
plt.xlabel("Execution Time (sec)")
plt.title("Prefix Sum Comparison")
plt.show()

```

</details>

#### 結果

![result_5.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3665424/6730a4ab-df8e-3b6f-35cb-b994b20688c5.png)

`sum_partial`を使う実装（`use_sp`）に比べて結構改善されていますね。内部メソッドに依存しない形にできますし個人的には好みです。


<details><summary>おまけ：区間加算・区間合計BITの改善全部載せ実装</summary>

```python
class BinaryIndexedTree:
    __slots__ = ("n", "tree")

    def __init__(self, n: int) -> None:
        self.n = n
        self.tree = [[0] * (self.n + 1), [0] * (self.n + 1)]

    @classmethod
    def from_array(cls, arr) -> "BinaryIndexedTree":
        ins = cls(len(arr))
        ins.tree[0][1:] = arr
        for i in range(1, ins.n):
            j = i + (i & -i)
            if j <= ins.n:
                ins.tree[0][j] += ins.tree[0][i]
        return ins

    def prefix_sum(self, i: int):
        res0 = res1 = 0
        tmp_i = i
        while tmp_i > 0:
            res0 += self.tree[0][tmp_i]
            res1 += self.tree[1][tmp_i]
            tmp_i -= tmp_i & -tmp_i
        return res0 + res1 * i

    def add_range(self, l: int, r: int, x) -> None:
        l += 1
        tmp_l, tmp_r = l, r
        while tmp_l <= self.n:
            self.tree[0][tmp_l] -= x * (l - 1)
            self.tree[1][tmp_l] += x
            tmp_l += tmp_l & -tmp_l
        while tmp_r <= self.n:
            self.tree[0][tmp_r] += x * r
            self.tree[1][tmp_r] -= x
            tmp_r += tmp_r & -tmp_r

    def add_point(self, i: int, x) -> None:
        """Add x at i (0-indexed)"""
        i += 1
        while i <= self.n:
            self.tree[0][i] += x
            i += i & -i

    def update_point(self, i: int, x) -> None:
        """Update i to x (0-indexed)"""
        self.add_point(i, x - self[i])

    def range_sum(self, l: int, r: int):
        """Range sum of [l, r) (0-indexed)"""
        res0 = res1 = 0
        tmp_l, tmp_r = l, r
        while tmp_l < tmp_r:
            res0 += self.tree[0][tmp_r]
            res1 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        res2 = res1
        while tmp_r < tmp_l:
            res0 -= self.tree[0][tmp_l]
            res1 -= self.tree[1][tmp_l]
            tmp_l -= tmp_l & -tmp_l
        while tmp_r > 0:
            res2 += self.tree[1][tmp_r]
            tmp_r -= tmp_r & -tmp_r
        return res0 + res1 * l + res2 * (r - l)

    def __getitem__(self, i: int):
        return self.range_sum(i, i + 1)

    def __iter__(self):
        yield from (self[i] for i in range(self.n))

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({list(self)})"

```

</details>

## 結び

というわけで`range_sum`の高速化の紹介と区間加算・区間合計BITへの適用の提案でした。実用上はどの実装でも困ることはないとは思いますが、こうした考えを通してデータ構造と仲良くなっておくと咄嗟の拡張にも対応できそうですね（願望）。
前回同様2Dへの拡張は未来の自分への宿題にします（この記事を読んでくださった方への遺題継承ともしたいです）。

この記事が参考になりましたら嬉しいです。よかったら左上のいいね押してね。