Segment Tree: Multiply Modulo in a Given Range

Introduction

Segment Tree is a data structure commonly used in computer science for efficient querying and updates of intervals or segments within an array. In this case, we will focus on solving a problem related to modular arithmetic.

Given an array of integers and two indices l and r, the task is to find the product of all elements in that range, modulo a given prime number p.

Solution

We can solve this problem efficiently using a Segment Tree. The idea is to represent each node in the tree by a range [L, R] and their product P. Then, we can query and update the tree recursively while keeping track of the modular multiplication.

Building the Tree

To build a Segment Tree for a list of integers arr, we can use a recursive approach. At each level of the tree, we divide the range in two halves until the range is reduced to a single element. Then, we assign its product to the node.

def build(start, end, node):
    if start == end:
        tree[node] = arr[start]
        return
    
    mid = (start + end) // 2
    build(start, mid, 2*node)
    build(mid+1, end, 2*node+1)
    
    tree[node] = (tree[2*node] * tree[2*node+1]) % p

Querying the Tree

To query for the product of a range [l, r], we also use a recursive approach. We check if the current node [L, R] is inside the range. If it is completely contained, we return its product. Otherwise, we split the query into the two halves that intersect with the range and multiply their results.

def query(left, right, start, end, node):
    if right < start or end < left:
        return 1
    
    if left <= start and end <= right:
        return tree[node]
    
    mid = (start + end) // 2
    p1 = query(left, right, start, mid, 2*node)
    p2 = query(left, right, mid+1, end, 2*node+1)
    
    return (p1 * p2) % p

Updating the Tree

To update the value of an element arr[i], we can also use a recursive approach. We find the node that represents the range that contains it and update its product. Then, we propagate the change upwards until the root node.

def update(idx, value, start, end, node):
    if start == end:
        arr[idx] = value
        tree[node] = value
        return
    
    mid = (start + end) // 2
    if start <= idx <= mid:
        update(idx, value, start, mid, 2*node)
    else:
        update(idx, value, mid+1, end, 2*node+1)
    
    tree[node] = (tree[2*node] * tree[2*node+1]) % p

Putting it all together

With these three functions, we can solve the problem of finding the product of a range modulo a given prime. The overall complexity is O(log n), where n is the size of the array, for each query and update.

arr = [...]
p = ...
tree = [0] * 4*len(arr)
build(0, len(arr)-1, 1)

# Query for product of range [l, r] modulo p
l, r = ..., ...
result = query(l, r, 0, len(arr)-1, 1)

# Update element i to value v
i, v = ..., ...
update(i, v, 0, len(arr)-1, 1)

Conclusion

Segment Trees are a powerful data structure that can be used to efficiently solve a wide range of problems related to intervals or segments in an array. In this case, we have seen how to use it to solve a problem related to modular arithmetic. With the functions presented here, it should be easy to implement and use a Segment Tree in your own code.