LoginSignup
45
27

More than 3 years have passed since last update.

@dn6049949

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

Last updated at Posted at 2020-04-05

はじめに

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

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$ の結果に相当します。


なおモノイドとは次の条件を満たす集合 $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の説明を行います。


まず、愚直な実装を考えると次のようになります。

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)$となっているため遅いです。


ここで次の図のようなデータ構造を考えます。( 結論から言うとこれが他ならぬSegment Treeです。)
segtree_1.jpg
※ []内の数字はSegment Treeにおけるインデックス、太字は入っているデータの値を示しています。


このSegment Treeを用いることで、 $update(i, x)$ の計算量は $O(logN)$ と増えますが、$query(l, r)$ の計算量も $O(logN)$ まで落とすことができます。

例を挙げると、例えば $query(1, 7)$ については下図の灰色の部分の和を求めればよく、愚直にやると5回必要な足し算が、たったの3回で済みます。
segtree_2.jpg

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


では実装(再帰型)について説明します。

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)$ の実装は次のとおりです。

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)$に完全に収まるときにその値を加算すればよいです。よって再帰を用いて以下のように書けます。

# 再帰関数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を実装すると次のとおりとなります。なおこれは一般のモノイドに対する実装です。

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)


非再帰型Segment Treeの実装

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

再帰のときと同じ構造でも実装できますが、インデックスを1ずらして1-indexedにすることで実装が非常に楽になります。1-indexedにした場合、下図のような添え字となります。
segtree_3.jpg

1-indexedにすることで、あるノード $i$ について次の表のような関係が得られます。この関係が実装の上で便利です。
親のインデックス $i/2$
左側の子のインデックス $2i$
右側の子のインデックス $2i+1$
$a_i$の値が代入されているインデックス $i+N$


update(i, x)の実装

$update(i, x)$ の実装は、再帰型のときとほぼ同じで、上記の関係に気を付けるだけです。実装例は次のとおりです。

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]


query(l, r)の実装

$query(l, r)$の実装ですが、ここで、Segment Treeのあるデータを加算する必要十分条件について考えると、次のとおりであることが分かります。証明は省きます。
「そのノードが表す区間が$[l,r)$に完全に収まっており、その親ノードが表す区間が$[l,r)$に完全に収まっていないとき」

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

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

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

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


非再帰型Segment Treeの実装のまとめ

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

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)
・AOJ DSL_2_B - Range Sum Query (RSQ)
・AtCoder ABC125 C - GCD on Blackboard

提出URL
【再帰】
・AOJ DSL_2_A
・AOJ DSL_2_B
・AtCoder ABC125 C

【非再帰】
・AOJ DSL_2_A
・AOJ DSL_2_B
・AtCoder ABC125 C

最後に

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

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

45
27
4

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
45
27