Introduction to Segment Tree

and program to construct a segment tree in Rust

Introduction

The Segment Tree is a commonly used tree like data structure that is used for range queries or subarray queries, like minimum element or sum of elements in a given subarray.

By using the Segment Tree, we can reduce the time taken for each range query to O( log n ), that takes O( n ) otherwise, but uses O( n ) auxiliary space.

The example of range queries on a given array are :

Sum/Product of integers in a given range.
Bitwise OR/AND of integers in given range.
Minimum/Maximum in a given range.
LCM/GCD in a given range.

If you got any of these queries to perform in Logarithmic time, segment tree is a good option.

Size of Segment Tree

For a given array containing n integers, size of segment tree is

2n-1 : If n is a power of 2
2k-1 : If n is not the power of 2, and k is the smallest power of 2 greater than n.

So, for example, if n = 9 => k = 16, and hence, size of segment tree = 31.

Also, it is not hard to see that space complexity of a segment tree = O( n )

Here is the function to calculate the height of a Segment Tree in Rust

fn calc_size(n:usize)->usize{
    // If n is a power of 2, return 2n-1
    if n.count_ones() == 1 {
        // Left shift is used to multiply n by 2.
        return (n<<1)-1;
    }

    // k denotes the height of the tree.

    // If n is not the power of 2, we have to take 2*k - 1
    // k can be calculated by 2^(position of MSB+1)
    // For example, for n = 7, position of MSB = 2, so k = 8
    // Hence k = 2^(64 - leading zeroes)

    let k = 1<<(64 - n.leading_zeros());
    return (k<<1)-1;
}

Structure of Segment Tree

The Segment Trees have a Binary Tree like structure, where each node has at most two children, each representing the left and right half of parent range respectively.

For Example : If a node represents the N elements from [0..N-1] ( inclusive ), then its Left Child will represent the elements from [0..N/2-1] and right child will represent the elements from [N/2..N-1].

Hence, we keep on dividing each node, until the node is representing a single element only.

Also, we store the elements in the array form. So, for ith node, its

Left Node is represented by element at index 2*i +1
Right Node is represented by element at index 2*i +2

Constructing a segment tree

Before we can perform our Range Queries, we have to construct a Segment Tree from the given array.

Algorithm

Calculate the size of segment tree and initialize an empty array/vector for segment tree with default value.
Initially, set low to 0, high to n-1 and position to 0, where size of original array is n.
If high is equal to low, then the current range represent only 1 element, and hence update it on the given position in the segment tree, and return.
If high is greater than low, then the current range represent 2 or more elements. So, we divide the range into 2 halves.
Recursively follow the steps 3 and 4 for left and right half of the range, for position 2×pos+1 and 2×pos+2 respectively.
Now, set the node at position to result of desired operator of left and right child.

Function

The program using above algorithm is given below. I used Multiplication operator for the reference.

All you have to do is call construct_segment_tree() function, and pass the reference to the original array and n, the size of the array.

Note : To change the operator, you can change the last line of construct_segment_tree_util() function.

fn calc_size(n:usize) ->usize{
    if n.count_ones() == 1 { return (n<<1)-1; }
    let k = 1<<(64 - n.leading_zeros());
    return (k<<1)-1;
}

// I am taking Vector of unsigned integers.
// You can take it as per your requirement
fn construct_segment_tree(arr:&Vec<usize>) -> Vec<usize> {
    // Calculate the size of the segment tree array first
    let n = arr.len();
    let s = calc_size(n);

    // Initialize the segment tree as an array of size s, and all elements are 0
    let mut segment_tree = vec![0 as usize; s];

    // Now, call the construct_segment_tree_util function to build the tree
    // Note : we use n-1 here, because we are using inclusive range.
    construct_segment_tree_util(arr, &mut segment_tree, 0, n-1, 0);

    return segment_tree;
}

fn construct_segment_tree_util(arr:&Vec<usize>, tree:&mut Vec<usize>, low:usize, high:usize, pos:usize){
    // The range is empty
    if high<low { return; }

    // Only 1 element in the range.
    // So, this is right position for the element
    if high == low {
        tree[pos] = arr[low];
        return;
    }

    // Now, multiple elements exist in the given range.
    // So, recursively perform this function for left and right half of range
    let mid = (high+low)/2;

    construct_segment_tree_util(arr, tree, low, mid, pos*2+1);
    construct_segment_tree_util(arr, tree, mid+1, high, pos*2+2);

    // Here, you can change operation as per requirement
    // I am using multiplication operator for reference
    tree[pos] = tree[pos*2+1]*tree[pos*2+2];
}

Time Complexity : O( n )
Space Complexity : O( n )

Conclusion

Segment tree is used frequently for Range Queries or Subarray Queries, especially in competitive programming.

It reduces the time complexity for each range query to O( log n ) from O( n ), but uses O( n ) auxiliary space. So, when you have large number of range queries, you can easily use Segment tree to save time.

In this article, we saw the code to construct a Segment Tree in Rust. We will see Range Queries in the next article. Here is the optimized function for easy access

fn cs(n:usize) ->usize{
    if n.count_ones() == 1 { return (n<<1)-1; }
    let k = 1<<(64 - n.leading_zeros());
    return (k<<1)-1;
}

fn cons_st(arr:&Vec<usize>) -> Vec<usize> {
    let n = arr.len();
    let s = cs(n);
    let mut segment_tree = vec![0 as usize; s];
    cons_st_util(arr, &mut segment_tree, 0, n-1, 0);
    return segment_tree;
}

fn cons_st_util(arr:&Vec<usize>, tree:&mut Vec<usize>, l:usize, h:usize, pos:usize){
    if h<l { return; }
    if h == l { tree[pos] = arr[l];return; }
    let mid = (h+l)/2;
    cons_st_util(arr, tree, l, mid, pos*2+1);
    cons_st_util(arr, tree, mid+1, h, pos*2+2);

    // Here, you can change operation as per requirement
    // I am using product operator for reference
    tree[pos] = tree[pos*2+1]*tree[pos*2+2];
}

Thank You