segmenttree creates a segment tree with a given array and function allowing queries to be done later in logn time function takes 2 values and returns a same type value import operator numarr segmenttree2 1 5 3 4 operator add tuplenumarr traverse doctest normalizewhitespace segmenttreenodestart0 end4 val15 segmenttreenodestart0 end2 val8 segmenttreenodestart3 end4 val7 segmenttreenodestart0 end1 val3 segmenttreenodestart2 end2 val5 segmenttreenodestart3 end3 val3 segmenttreenodestart4 end4 val4 segmenttreenodestart0 end0 val2 segmenttreenodestart1 end1 val1 numarr update1 5 tuplenumarr traverse doctest normalizewhitespace segmenttreenodestart0 end4 val19 segmenttreenodestart0 end2 val12 segmenttreenodestart3 end4 val7 segmenttreenodestart0 end1 val7 segmenttreenodestart2 end2 val5 segmenttreenodestart3 end3 val3 segmenttreenodestart4 end4 val4 segmenttreenodestart0 end0 val2 segmenttreenodestart1 end1 val5 numarr queryrange3 4 7 numarr queryrange2 2 5 numarr queryrange1 3 13 maxarr segmenttree2 1 5 3 4 max for node in maxarr traverse printnode segmenttreenodestart0 end4 val5 segmenttreenodestart0 end2 val5 segmenttreenodestart3 end4 val4 segmenttreenodestart0 end1 val2 segmenttreenodestart2 end2 val5 segmenttreenodestart3 end3 val3 segmenttreenodestart4 end4 val4 segmenttreenodestart0 end0 val2 segmenttreenodestart1 end1 val1 maxarr update1 5 for node in maxarr traverse printnode segmenttreenodestart0 end4 val5 segmenttreenodestart0 end2 val5 segmenttreenodestart3 end4 val4 segmenttreenodestart0 end1 val5 segmenttreenodestart2 end2 val5 segmenttreenodestart3 end3 val3 segmenttreenodestart4 end4 val4 segmenttreenodestart0 end0 val2 segmenttreenodestart1 end1 val5 maxarr queryrange3 4 4 maxarr queryrange2 2 5 maxarr queryrange1 3 5 minarr segmenttree2 1 5 3 4 min for node in minarr traverse printnode segmenttreenodestart0 end4 val1 segmenttreenodestart0 end2 val1 segmenttreenodestart3 end4 val3 segmenttreenodestart0 end1 val1 segmenttreenodestart2 end2 val5 segmenttreenodestart3 end3 val3 segmenttreenodestart4 end4 val4 segmenttreenodestart0 end0 val2 segmenttreenodestart1 end1 val1 minarr update1 5 for node in minarr traverse printnode segmenttreenodestart0 end4 val2 segmenttreenodestart0 end2 val2 segmenttreenodestart3 end4 val3 segmenttreenodestart0 end1 val2 segmenttreenodestart2 end2 val5 segmenttreenodestart3 end3 val3 segmenttreenodestart4 end4 val4 segmenttreenodestart0 end0 val2 segmenttreenodestart1 end1 val5 minarr queryrange3 4 3 minarr queryrange2 2 5 minarr queryrange1 3 3 update an element in logn time param i position to be update param val new value import operator numarr segmenttree2 1 5 3 4 operator add numarr update1 5 numarr queryrange1 3 13 get range query value in logn time param i left element index param j right element index return element combined in the range i j import operator numarr segmenttree2 1 5 3 4 operator add numarr update1 5 numarr queryrange3 4 7 numarr queryrange2 2 5 numarr queryrange1 3 13 range in left child tree range in left child tree and right child tree range in right child tree import operator num_arr segmenttree 2 1 5 3 4 operator add tuple num_arr traverse doctest normalize_whitespace segmenttreenode start 0 end 4 val 15 segmenttreenode start 0 end 2 val 8 segmenttreenode start 3 end 4 val 7 segmenttreenode start 0 end 1 val 3 segmenttreenode start 2 end 2 val 5 segmenttreenode start 3 end 3 val 3 segmenttreenode start 4 end 4 val 4 segmenttreenode start 0 end 0 val 2 segmenttreenode start 1 end 1 val 1 num_arr update 1 5 tuple num_arr traverse doctest normalize_whitespace segmenttreenode start 0 end 4 val 19 segmenttreenode start 0 end 2 val 12 segmenttreenode start 3 end 4 val 7 segmenttreenode start 0 end 1 val 7 segmenttreenode start 2 end 2 val 5 segmenttreenode start 3 end 3 val 3 segmenttreenode start 4 end 4 val 4 segmenttreenode start 0 end 0 val 2 segmenttreenode start 1 end 1 val 5 num_arr query_range 3 4 7 num_arr query_range 2 2 5 num_arr query_range 1 3 13 max_arr segmenttree 2 1 5 3 4 max for node in max_arr traverse print node segmenttreenode start 0 end 4 val 5 segmenttreenode start 0 end 2 val 5 segmenttreenode start 3 end 4 val 4 segmenttreenode start 0 end 1 val 2 segmenttreenode start 2 end 2 val 5 segmenttreenode start 3 end 3 val 3 segmenttreenode start 4 end 4 val 4 segmenttreenode start 0 end 0 val 2 segmenttreenode start 1 end 1 val 1 max_arr update 1 5 for node in max_arr traverse print node segmenttreenode start 0 end 4 val 5 segmenttreenode start 0 end 2 val 5 segmenttreenode start 3 end 4 val 4 segmenttreenode start 0 end 1 val 5 segmenttreenode start 2 end 2 val 5 segmenttreenode start 3 end 3 val 3 segmenttreenode start 4 end 4 val 4 segmenttreenode start 0 end 0 val 2 segmenttreenode start 1 end 1 val 5 max_arr query_range 3 4 4 max_arr query_range 2 2 5 max_arr query_range 1 3 5 min_arr segmenttree 2 1 5 3 4 min for node in min_arr traverse print node segmenttreenode start 0 end 4 val 1 segmenttreenode start 0 end 2 val 1 segmenttreenode start 3 end 4 val 3 segmenttreenode start 0 end 1 val 1 segmenttreenode start 2 end 2 val 5 segmenttreenode start 3 end 3 val 3 segmenttreenode start 4 end 4 val 4 segmenttreenode start 0 end 0 val 2 segmenttreenode start 1 end 1 val 1 min_arr update 1 5 for node in min_arr traverse print node segmenttreenode start 0 end 4 val 2 segmenttreenode start 0 end 2 val 2 segmenttreenode start 3 end 4 val 3 segmenttreenode start 0 end 1 val 2 segmenttreenode start 2 end 2 val 5 segmenttreenode start 3 end 3 val 3 segmenttreenode start 4 end 4 val 4 segmenttreenode start 0 end 0 val 2 segmenttreenode start 1 end 1 val 5 min_arr query_range 3 4 3 min_arr query_range 2 2 5 min_arr query_range 1 3 3 update an element in log n time param i position to be update param val new value import operator num_arr segmenttree 2 1 5 3 4 operator add num_arr update 1 5 num_arr query_range 1 3 13 get range query value in log n time param i left element index param j right element index return element combined in the range i j import operator num_arr segmenttree 2 1 5 3 4 operator add num_arr update 1 5 num_arr query_range 3 4 7 num_arr query_range 2 2 5 num_arr query_range 1 3 13 range in left child tree range in left child tree and right child tree range in right child tree 7 5 13
from collections.abc import Sequence
from queue import Queue


class SegmentTreeNode:
    def __init__(self, start, end, val, left=None, right=None):
        self.start = start
        self.end = end
        self.val = val
        self.mid = (start + end) // 2
        self.left = left
        self.right = right

    def __repr__(self):
        return f"SegmentTreeNode(start={self.start}, end={self.end}, val={self.val})"


class SegmentTree:
    

    def __init__(self, collection: Sequence, function):
        self.collection = collection
        self.fn = function
        if self.collection:
            self.root = self._build_tree(0, len(collection) - 1)

    def update(self, i, val):
        
        self._update_tree(self.root, i, val)

    def query_range(self, i, j):
        
        return self._query_range(self.root, i, j)

    def _build_tree(self, start, end):
        if start == end:
            return SegmentTreeNode(start, end, self.collection[start])
        mid = (start + end) // 2
        left = self._build_tree(start, mid)
        right = self._build_tree(mid + 1, end)
        return SegmentTreeNode(start, end, self.fn(left.val, right.val), left, right)

    def _update_tree(self, node, i, val):
        if node.start == i and node.end == i:
            node.val = val
            return
        if i <= node.mid:
            self._update_tree(node.left, i, val)
        else:
            self._update_tree(node.right, i, val)
        node.val = self.fn(node.left.val, node.right.val)

    def _query_range(self, node, i, j):
        if node.start == i and node.end == j:
            return node.val

        if i <= node.mid:
            if j <= node.mid:
                
                return self._query_range(node.left, i, j)
            else:
                
                return self.fn(
                    self._query_range(node.left, i, node.mid),
                    self._query_range(node.right, node.mid + 1, j),
                )
        else:
            
            return self._query_range(node.right, i, j)

    def traverse(self):
        if self.root is not None:
            queue = Queue()
            queue.put(self.root)
            while not queue.empty():
                node = queue.get()
                yield node

                if node.left is not None:
                    queue.put(node.left)

                if node.right is not None:
                    queue.put(node.right)


if __name__ == "__main__":
    import operator

    for fn in [operator.add, max, min]:
        print("*" * 50)
        arr = SegmentTree([2, 1, 5, 3, 4], fn)
        for node in arr.traverse():
            print(node)
        print()

        arr.update(1, 5)
        for node in arr.traverse():
            print(node)
        print()

        print(arr.query_range(3, 4))  
        print(arr.query_range(2, 2))  
        print(arr.query_range(1, 3))  
        print()