非再帰型Segment TreeのPythonによる実装

# はじめに
タイトルのとおり、この記事では**Python**による**非再帰型Segment Tree**の実装を紹介したいと思います。
前提知識を「ほぼ」[^1] 必要としないように**Segment Tree**の説明から入るので、もう知ってるという方は読み飛ばしてください。→[非再帰型Segment Treeの実装](#非再帰型segment-treeの実装)

[^1]: 計算量の概念や、Pythonの記法などについては理解しておく必要があります。

# Segment Treeとは
通称**セグ木**、**セグメント木**などと呼ばれているデータ構造で、長さ $N$ の配列{$a_i$}$_{i=0}^{N-1}$に対して、**モノイド** $(S,•)$ に関する次の2つの操作がどちらも時間計算量が $O(logN)$ で行えます。

* $update(i, x)$:　 $a_i$ を $x$ に変更する
* $query(l, r)$:　 半開区間 $[l,r)$ に対して $a_l•a_{l+1}•$…$•a_{r-1}$ を返す

**モノイド**として、一般的な足し算を考えると、$query(l, r)$ は $a_l$ から $a_{r-1}$ までの区間和、つまり $\Sigma_{i=l}^{r-1}a_i$ の結果に相当します。
<br>
なお**モノイド**とは次の条件を満たす集合 $S$ と二項演算子 $•$ の組 $(S,•)$ のことです。

* 任意の $a, b, c \in S$ に対して、$(a•b)•c=a•(b•c)$ (**結合律**)
* ある $e \in S$ が存在して、任意の $a \in S$ に対して、$e•a=a•e=a$ (**単位元の存在**)

**モノイド**の例としては、$(\mathbb{N},+)$、$(\mathbb{Z},-)$、$(\mathbb{R},min)$ や $(\mathbb{Q},max)$ などが上げられます。

以下では**モノイド**として $(\mathbb{R},+)$ を例に挙げ、**Segment Tree**の説明を行います。
<br>
まず、愚直な実装を考えると次のようになります。

```python
def update(i, x):
    a[i] = x

def query(l, r):
    res = 0
    for i in range(l, r):
        res += a[i]
    return res
```

これらの時間計算量は $update(i, x)$ が $O(1)$ と高速ですが、$query(l, r)$ が $O(N)$となっているため遅いです。
<br>
ここで次の図のようなデータ構造を考えます。( 結論から言うとこれが他ならぬ**Segment Tree**です。)
![segtree_1.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/614181/bb7a9041-c44f-3d69-5d38-95e44222c53c.jpeg)
※ []内の数字は**Segment Tree**におけるインデックス、太字は入っているデータの値を示しています。
<br>
この**Segment Tree**を用いることで、 $update(i, x)$ の計算量は $O(logN)$ と増えますが、$query(l, r)$ の計算量も $O(logN)$ まで落とすことができます。

例を挙げると、例えば $query(1, 7)$ については下図の灰色の部分の和を求めればよく、愚直にやると5回必要な足し算が、たったの3回で済みます。
![segtree_2.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/614181/e19c7272-d937-2430-b10e-8c86460af659.jpeg)

実装についてメインに説明したいので詳しい証明は行いませんが、あらゆる $l,r(l < r)$ の組み合わせに対して、$query(l, r)$ は $O(logN)$ で動作します。
<br>
では実装(**再帰型**)について説明します。


#### update(i, x)の実装
$update(i, x)$ について更新すべき値は $a_i$(上図の1番下の層), $a_i+a_j$(下から2番目の層), $a_i+a_j+a_k+a_l$(下から3番目の層), $a_i+a_j+a_k+a_l+$…(下から4番目の層)
の計4つです。この数は**Segment Tree**を完全二分木とみなしたときの木の高さと同じなので$O(logN)$です。

またインデックスを上図のようにしておくと、更新すべき値を持つインデックスは順番に
$i+N-1,\ \lfloor \frac{(i+N-1)-1}{2} \rfloor,\ \lfloor \frac{\lfloor \frac{(i+N-1)-1}{2} \rfloor-1}{2} \rfloor,\ ...$
となっています。
これは、Segment Tree において $a_i$ の値が代入されているインデックスが $i+N-1$ であるため、また子 $x$ の親のインデックスが $\lfloor \frac{x-1}{2} \rfloor$ となるためです。

以上のことから、$update(i, x)$ の実装は次のとおりです。

```python
def update(i, x):
    i += N-1 # 1番下の層におけるインデックス
    dat[i] = x

    # 層をのぼりながら値を更新
    while i > 0:
        i = (i-1) >> 1 # 1つ上の層のインデックス(完全二分木における親)

        # 下の層2つの和を代入(完全二分木における子同士の和)
        dat[i] = dat[i*2+1] + dat[i*2+2]
```

#### query(l, r)の実装
$query(l, r)$は「あるデータに代入されている区間の値」が $[l, r)$に完全に収まるときにその値を加算すればよいです。よって**再帰**を用いて以下のように書けます。

```python
# 再帰関数query(l, r, k, L, R)
# L : 区間[L,R)の左端(初期値0), R : 区間[L,R)の右端(初期値N)
# k : 区間[L,R)に対する演算結果を保持しているインデックス(初期値0)

def query(l, r, k=0, L=0, R=None):
    # Rの初期化
    if R is None:
        R = N

    # もし区間[l,r)と区間[L,R)が交差しないならば0を返す
    if R <= l or r <= L:
        return 0

    # 区間[L,R)が区間[l,r)に完全に収まっていればその値を返す
    if l <= L and R <= r:
        return dat[k]

    # それ以外のとき、木を潜って2つの子を見る
    else:
        lres = query(l, r, k*2+1, L, (L+R) >> 1)
        rres = query(l, r, k*2+2, (L+R) >> 1, R)
        return lres + rres
```

#### Segment Tree(再帰型)の実装のまとめ
以上を踏まえてSegment Treeを実装すると次のとおりとなります。なおこれは一般のモノイドに対する実装です。

```Python
class SegmentTree:
    # 初期化処理
    # f : SegmentTreeにのせるモノイド
    # default : fに対する単位元
    def __init__(self, size, f=lambda x,y : x+y, default=0):
        self.size = 2**(size-1).bit_length() # 簡単のため要素数Nを2冪にする
        self.default = default
        self.dat = [default]*(self.size*2-1) # 要素を単位元で初期化
        self.f = f

    def update(self, i, x):
        i += self.size-1
        self.dat[i] = x
        while i > 0:
            i = (i-1) >> 1
            self.dat[i] = self.f(self.dat[i*2+1], self.dat[i*2+2])

    def query(self, l, r, k=0, L=0, R=None):
        if R is None:
            R = self.size
        if R <= l or r <= L:
            return self.default
        if l <= L and R <= r:
            return self.dat[k]
        else:
            lres = self.query(l, r, k*2+1, L, (L+R) >> 1)
            rres = self.query(l, r, k*2+2, (L+R) >> 1, R)
            return self.f(lres, rres)
```
<br>

# 非再帰型Segment Treeの実装
さて、ここまではやや長い前置きで、ここからが本題です。上記の**SegmentTree**では**再帰関数**を用いているため、やや遅いです。そこで**非再帰関数**による実装が有用となってきます。

再帰のときと同じ構造でも実装できますが、インデックスを1ずらして**1-indexed**にすることで実装が非常に楽になります。**1-indexed**にした場合、下図のような添え字となります。
![segtree_3.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/614181/2954cf8e-9630-d07a-2ed1-fefc4f430831.jpeg)

**1-indexed**にすることで、あるノード $i$ について次の表のような関係が得られます。この関係が実装の上で便利です。


|||
|---|---|
|親のインデックス|$i/2$|
|左側の子のインデックス|$2i$|
|右側の子のインデックス|$2i+1$|
|$a_i$の値が代入されているインデックス|$i+N$|
<br>
#### update(i, x)の実装
$update(i, x)$ の実装は、再帰型のときとほぼ同じで、上記の関係に気を付けるだけです。実装例は次のとおりです。

```Python
def update(i, x):
    i += N # 1番下の層におけるインデックス
    dat[i] = x

    # 層をのぼりながら値を更新
    while i > 0:
        i >>= 1 # 1つ上の層のインデックス(完全二分木における親)
        # 下の層2つの和を代入(完全二分木における子同士の和)
        dat[i] = dat[i*2] + dat[i*2+1]
```
<br>
#### query(l, r)の実装
$query(l, r)$の実装ですが、ここで、**Segment Tree**のあるデータを加算する必要十分条件について考えると、次のとおりであることが分かります。証明は省きます。
「**そのノードが表す区間が$[l,r)$に完全に収まっており、その親ノードが表す区間が$[l,r)$に完全に収まっていないとき**」

あるノードについて、表す区間が$[l,r)$に収まっており、親ノードの表す区間が**左**方向にはみ出すとき、そのノードは**右**側の子であるため、**奇数**のインデックスを持ちます。
反対に**右**方向にはみ出す時、そのノードは**左**側の子であるため、**偶数**のインデックスを持ちます。

よって、左側については木を下からのぼっていき、親の区間がはみ出すならば(現在見ているノードのインデックスが**奇数**ならば)、そのノードの値を加算し、インデックスを1つ右にずらすということを繰り返せばよいです。
右側についても同様で、親の区間がはみ出すときに加算してインデックスを1つ左にずらせばよいです。ただし、**半開区間**で見ているので、はみ出すかどうかの判定は現在見ているノードのインデックスが**奇数**であるかどうかで行います。また加算するのは**1つ左**のノードであることに注意が必要です。

実装例は次のとおりです。

```Python
def query(l, r):
    # 1番下の層におけるインデックスに初期化
    l += N
    r += N

    # 左側の答えと右側の答えを0で初期化
    lres, rres = 0, 0

    # lとrが重なるまで上記の判定を用いて加算を実行
    while l < r:
        # lが奇数ならば、dat[l]を加算
        if l & 1:
            lres += dat[l]
            l += 1

        # rが奇数ならば、dat[r-1]を加算
        if r & 1:
            r -= 1
            rres += dat[r]

        # 木をのぼる
        l >>= 1
        r >>= 1

    res = lres + rres
    return res
```
<br>
#### 非再帰型Segment Treeの実装のまとめ
以上のことを踏まえて、非再帰型Segment Treeを実装すると次のとおりとなります。なおこれは一般のモノイドに対する実装です。

```Python
class SegmentTree:
    # 初期化処理
    # f : SegmentTreeにのせるモノイド
    # default : fに対する単位元
    def __init__(self, size, f=lambda x,y : x+y, default=0):
        self.size = 2**(size-1).bit_length() # 簡単のため要素数Nを2冪にする
        self.default = default
        self.dat = [default]*(self.size*2) # 要素を単位元で初期化
        self.f = f

    def update(self, i, x):
        i += self.size
        self.dat[i] = x
        while i > 0:
            i >>= 1
            self.dat[i] = self.f(self.dat[i*2], self.dat[i*2+1])

    def query(self, l, r):
        l += self.size
        r += self.size
        lres, rres = self.default, self.default
        while l < r:
            if l & 1:
                lres = self.f(lres, self.dat[l])
                l += 1

            if r & 1:
                r -= 1
                rres = self.f(self.dat[r], rres) # モノイドでは可換律は保証されていないので演算の方向に注意
            l >>= 1
            r >>= 1
        res = self.f(lres, rres)
        return res
```

# 比較
**Segment Tree**を使って解くことができる問題を使って実行時間の比較を行いました。結果は次の表のとおりで、非再帰のほうが短い時間で実行できることが確認できました。


|問題名|再帰|非再帰|
|---|---|---|
|AOJ DSL_2_A|AC (**3.08 s**)|AC (**1.73 s**)|
|AOJ DSL_2_B|AC (**3.73 s**)|AC (**1.78 s**)|
|AtCoder ABC125 C|AC (**963 ms**)|AC (**462 ms**)|

**問題URL**
・[AOJ DSL_2_A - Range Minimum Query (RMQ)](http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_A)
・[AOJ DSL_2_B - Range Sum Query (RSQ)](http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_B)
・[AtCoder ABC125 C - GCD on Blackboard](https://atcoder.jp/contests/abc125/tasks/abc125_c)


**提出URL**
【再帰】
・[AOJ DSL_2_A](http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326213#1)
・[AOJ DSL_2_B](http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326207#1)
・[AtCoder ABC125 C](https://atcoder.jp/contests/abc125/submissions/17120641)

【非再帰】
・[AOJ DSL_2_A](http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326214#1)
・[AOJ DSL_2_B](http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326206#1)
・[AtCoder ABC125 C](https://atcoder.jp/contests/abc125/submissions/17120647)


# 最後に
Qiitaどころか、生まれて初めて記事を書いたので至らぬことだらけだと思います。「ここがおかしい」「変なことが書いてある」等ございましたら、ご報告お願いします。
ありがとうございました。