Python implementation of non-recursive Segment Tree

Introduction

As the title suggests, this article would like to introduce you to the ** non-recursive Segment Tree ** implementation with ** Python **. I will start from the explanation of ** Segment Tree ** so that I do not need the prerequisite knowledge "almost" [^ 1], so if you already know it, please skip it. → [Implementation of non-recursive Segment Tree](# Implementation of non-recursive segment-tree)

[^ 1]: You need to understand the concept of computational complexity and Python notation.

What is Segment Tree?

A data structure commonly known as ** segment tree **, ** segment tree **, etc., an array of length $ N $ {$ a_i $} $ _ {i = 0} ^ {N-1} $ On the other hand, the following two operations related to ** monoid ** $ • $ can be performed with the time complexity of $ O (logN) $.

• $ update (i, x) $: Change $ a_i $ to $ x $
• $ query (l, r) $: Returns $ a_l • a_ {l + 1} • $… $ • a_ {r-1} $ for the half-open interval $ [l, r) $

As a ** monoid **, considering general addition, $ query (l, r) $ is the sum of the intervals from $ a_l $ to $ a_ {r-1} $, that is, $ \ Sigma_ {i = l} Corresponds to the result of ^ {r-1} a_i $.
Note that ** monoid ** is an operator $ • $ that satisfies the following conditions.

• For any $ a, b, c $, $ (a • b) • c = a • (b • c) $ (** associative law **)
• Some $ e $ exists, and for any $ a $, $ e • a = a • e = a $ (** identity element exists **)

Examples of ** monoids ** include addition and multiplication, $ min $ and $ max $.

In the following, ** Segment Tree ** will be explained using addition $ + $ as an example as ** monoid **.
First, considering a straightforward implementation:

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

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

These time complexitys are fast, with $ update (i, x) $ being $ O (1) $, but slow because $ query (l, r) $ is $ O (N) $.
Now consider the data structure shown in the following figure. (In conclusion, this is none other than ** Segment Tree **.)

• The numbers in [] indicate the index in ** Segment Tree **, and the bold letters indicate the values of the contained data.
  By using this ** Segment Tree **, the complexity of $ update (i, x) $ increases to $ O (logN) $, but the complexity of $ query (l, r) $ is also $ O ( You can drop it down to logN) $.

For example, for $ query (1, 7) $, you can find the sum of the gray parts in the figure below, and if you do it honestly, you only need to add 5 times, but only 3 times.

I won't give a detailed proof because I want to explain the implementation mainly, but for any combination of $ l, r (l <r) $, $ query (l, r) $ works with $ O (logN) $. I will.
Now let's talk about the implementation (** recursive **). For the sake of simplicity, we will fix it as $ N = 8 $, but the same applies to cases other than $ 8 $.

Implementation of update (i, x)

For $ update (i, x) $, the values to be updated are $ a_i $ (bottom layer in the figure above), $ a_i + a_j $ (second layer from bottom), $ a_i + a_j + a_k + a_l $ (3rd layer from the bottom), $ a_i + a_j + a_k + a_l + $… (4th layer from the bottom) There are a total of four. This number is $ O (logN) $ because it is the same as the height of the tree when ** Segment Tree ** is regarded as a complete binary tree.

Also, if the indexes are set as shown above, the indexes with the values to be updated will be in order. i+N-1, ((i+N-1)-1)/2, (((i+N-1)-1)/2-1)/2, ((((i+N-1)-1)/2-1)/2-1)/2 It has become. ($ / $ Is rounded down and divided.) I won't explain it in detail here either, but you can see that ** Segment Tree ** is a complete binary tree.

From the above, the implementation of $ update (i, x) $ is as follows.

def update(i, x):
    i += N-1 #Index in the bottom layer
    dat[i] = x

    #Update values while climbing layers
    while i > 0:
        i = (i-1) >> 1 #Index of the next higher layer(Parent in a complete binary tree)

        #Substitute the sum of the two lower layers(Sum of children in a complete binary tree)
        dat[i] = dat[i*2+1] + dat[i*2+2]

Implementation of query (l, r)

$ query (l, r) $ should be added when the "value of the interval assigned to a certain data" completely fits in $ [l, r) $. Therefore, you can write as follows using ** recursion **.

#Recursive function query(l, r, k, L, R)
# L :section[L,R)Left edge of(Initial value 0), R :section[L,R)Right end of(Initial value N)
# k :section[L,R)Index that holds the operation result for(Initial value 0)

def query(l, r, k=0, L=0, R=None):
    #Initialization of R
    if R is None:
        R = N

    #If section[l,r)And section[L,R)Returns 0 if does not intersect
    if R <= l or r <= L:
        return 0

    #section[L,R)Is the section[l,r)Returns that value if it fits perfectly in
    if l <= L and R <= r:
        return dat[k]

    #At other times, dive into the tree and see the two children
    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

Summary of Segment Tree (recursive) implementation

Based on the above, the Segment Tree is implemented as follows. This is an implementation for general monoids.

class SegmentTree:
    #Initialization process
    # f :Monoids on the Segment Tree
    # default :Identity element for f
    def __init__(self, size, f=lambda x,y : x+y, default=0):
        self.size = 2**(size-1).bit_length() #Set the number of elements N to 2 to the power for simplicity
        self.default = default
        self.dat = [default]*(self.size*2-1) #Initialize elements in identity element
        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)

Implementation of non-recursive Segment Tree

By the way, this is a rather long introduction, and the main subject is from here. The above ** SegmentTree ** uses a ** recursive function **, so it's a bit slow. Therefore, implementation by ** non-recursive function ** becomes useful.

It can be implemented with the same structure as for recursion, but it is much easier to implement by shifting the index by 1 to ** 1-indexed **. If **1-indexed ** is set, the subscript will be as shown in the figure below.

By making it ** 1-indexed **, you can get the relationship shown in the following table for a node $ i $. This relationship is useful for implementation.

def update(i, x):
    i += N #Index in the bottom layer
    dat[i] = x

    #Update values while climbing layers
    while i > 0:
        i >>= 1 #Index of the next higher layer(Parent in a complete binary tree)
        #Substitute the sum of the two lower layers(Sum of children in a complete binary tree)
        dat[i] = dat[i*2] + dat[i*2+1]

For a node, the interval represented is within $ [l, r) $, and when the interval represented by the parent node extends in the ** left ** direction, the node is a child on the ** right ** side. It has an index of ** odd **. On the other hand, when it extends in the ** right ** direction, the node is a child on the ** left ** side, so it has an index of ** even **.

Therefore, for the left side, climb the tree from the bottom, and if the parent section extends beyond (if the index of the node you are currently viewing is ** odd **), add the value of that node and add one index. You just have to repeat the process of shifting to the right. The same applies to the right side, and when the parent section extends beyond, it can be added and the index can be shifted to the left by one. However, since we are looking at the ** half-open interval **, we will judge whether or not it will extend based on whether the index of the node we are currently looking at is ** odd **. Also note that it is the node ** 1 left ** that adds.

The implementation example is as follows.

def query(l, r):
    #Initialize to index in the bottom layer
    l += N
    r += N

    #Initialize the answer on the left and the answer on the right with 0
    lres, rres = 0, 0

    #Perform addition using the above judgment until l and r overlap
    while l < r:
        #If l is odd, dat[l]Add
        if l & 1:
            lres += dat[l]
            l += 1

        #If r is odd, dat[r-1]Add
        if r & 1:
            r -= 1
            rres += dat[r]

        #Climb a tree
        l >>= 1
        r >>= 1

    res = lres + rres
    return res

class SegmentTree:
    #Initialization process
    # f :Monoids on the Segment Tree
    # default :Identity element for f
    def __init__(self, size, f=lambda x,y : x+y, default=0):
        self.size = 2**(size-1).bit_length() #Set the number of elements N to 2 to the power for simplicity
        self.default = default
        self.dat = [default]*(self.size*2) #Initialize elements in identity element
        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) #Commutative law is not guaranteed for monoids, so be careful about the direction of calculation.
            l >>= 1
            r >>= 1
        res = self.f(lres, rres)
        return res

Comparison

We compared execution times using problems that can be solved using ** Segment Tree **. The results are shown in the following table, confirming that non-recursive can be executed in a shorter time.

** Problem URL ** AOJ DSL_2_A http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_A

AOJ DSL_2_B http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_B

AtCoder ABC_125_C https://atcoder.jp/contests/abc125/tasks/abc125_c

** Submission URL ** AOJ DSL_2_A (recursive) http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326213#1

AOJ DSL_2_A (non-recursive) http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326214#1

AOJ DSL_2_B (recursive) http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326207#1

AOJ DSL_2_B (non-recursive) http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=4326206#1

AtCoder ABC_125_C (recursive) https://atcoder.jp/contests/abc125/submissions/11596118

AtCoder ABC_125_C (non-recursive) https://atcoder.jp/contests/abc125/submissions/11596126

Finally

Far from Qiita, I wrote an article for the first time in my life, so I think it's full of things. If you have any questions such as "this is strange" or "something strange is written", please report it. Thank you very much.