What is a segment tree used for?

A segment tree is a binary tree in which the information about a linear data structure segment, such as an array, is stored in the nodes. Mainly, it helps us resolve range queries.

Queries to Calculate the Sum of the Array Elements Consisting of an Odd Number of Divisors

Problem

Given an array A of size n and an array Query[][2] of size q consisting of queries of the form {L, R}, find the sum of array elements from the index range [L, R], having an odd number of divisors. Assume all the array elements are positive.

Input:

An integer n.
The next n lines contain the array elements.
The next line contains an integer q that represents the size of the Query array.
The next q lines contain two integers, L and R, representing the array's left and right indexes.

Output: For each query, print the answer in a new line.

The questions you can ask the interviewer:

What are the constraints on n and q?
What is the range of the elements of A?

Alert Ninjas: Don’t stop now! Enroll in the Competitive Programming course today. Learn the art of writing time and space-efficient code and compete in code competitions worldwide.

Example:

Input: n = 9, A[] = {1, 2, 3, 4, 5, 6, 7, 8, 9}, q = 3, Query[][] = {{0, 3}, {1, 5}, {4, 8}}

Output: 5

Explanation:

Query 1: Array elements from [0, 3] are {1, 2, 3, 4}. Out of them, only 1,4 has an odd number of divisors. Therefore, the sum is 5.

Query 2: Elements from indices [1, 5] are {2, 3, 4, 5, 6}. Only 4 have an odd number of divisors. Hence, the sum is 4.

Query 3: Elements from the indices [4, 8] are {5, 6, 7, 8, 9}. Out of them, only 9 have an odd number of divisors. Therefore, the sum is 9.

Recommended: Try to solve it yourself before moving on to the solution.

Solution

Naive Approach:

The most obvious solution that comes to our mind is that for each query, traverse the array A, over the range [L, R] and determine the sum of the array elements in the range [L, R] with odd numbers of divisors, then print the resultant sum.

The time complexity of the above approach is O( n * q * √n). As the time complexity to solve each query is O(q) *, the time complexity to traverse the index range [L, R] is O(n) *, the time complexity to find the divisors of each element in that range is O(√n).

Can we do better?

Segment Tree Approach:

Idea: The idea is to use Segment Trees.

Use the fact that the number of divisors of a positive integer is odd if only the integer is a perfect square. As a result, replace all the array elements that do not have the odd number of divisors with 0. Then, build a segment tree to answer each query.

Algorithm:

Traverse array A and replace the elements that are not a perfect square with 0.
Build a segment tree of the given array elements.
Iterate over the array Queries, and for each query, calculate the sum of the [L, R] range from the segment tree.
Print the answer in a separate line.

Time Complexity: O(q * log(n)); Time complexity of building the segment tree is O(logn) * time complexity of iterating over the array Queries.

Space Complexity: O(n); Space used to build the segment tree.

Let's try to reduce the time complexity further.

Better Approach:

Idea: The idea is to use Dynamic Programming such that dp[i] will be A[i] if A[i] is a perfect square, 0 otherwise. Find the prefix sum of the dp array to answer each query.

Algorithm:

Initialize a dp array of size n with values 0.
Iterate through array A using variable i and change dp[i] to A[i] if A[i] is a perfect square.
Calculate the prefix sum of the dp array.
Iterate over all the queries, and the answer for each query in the range [L, R] will be (dp[R] – dp[L – 1]).

C++

#include <bits/stdc++.h>
using namespace std;

void Sum(int n, int q, int A[], vector<vector<int> > &Query)
{
// Initialise the dp array with initial value 0.
vector<int> dp(n,0);

// Traverse the array A.
for (int i = 0; i < n; i++) {
int x = sqrt(A[i]);

// if A[i] is a perfect square, update DP[i] to A[i].
if (x * x == A[i])
dp[i] = A[i];
}

// Find the prefix sum of the dp array.
for (int i = 1; i < n; i++) {
dp[i] += dp[i - 1];
}

// Iterate over all the queries.
for (int i = 0; i < q; i++) {

int l = Query[i][0];
int r = Query[i][1];

// Find the sum for each query and print in a new line.
if (l == 0) {
cout << dp[r] <<endl;
}
else {
cout << dp[r] - dp[l - 1] <<endl;
}
}
}

int main()
{
int n = 9;
int A[] = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
int q = 3;
vector<vector<int> > Query = { {0, 3}, {1, 5}, {4, 8} };
Sum(n, q, A, Query);
return 0;
}

Input: n = 9, A[] = {1, 2, 3, 4, 5, 6, 7, 8, 9}, q = 3, Query[][] = {{0, 3}, {1, 5}, {4, 8}}.

Output:

Time Complexity: The time complexity is either O(n) or O(q), whichever is greater.

Space Complexity: O(n); Space used to create the DP array.

Also see, Euclid GCD Algorithm

FAQs

What is a segment tree used for?
A segment tree is a binary tree in which the information about a linear data structure segment, such as an array, is stored in the nodes. Mainly, it helps us resolve range queries.

Key takeaways

In this article, we designed an algorithm to solve all of the queries about calculating the sum of the elements of an array in a given range [L, R] with an odd number of divisors.

We solved it using three different approaches:

The naive approach in which for each query we were traversing the array, over the range [L, R] and calculating the sum of the array elements in the range [L, R] with odd numbers of divisors. It was taking O(n^3/2* q) time.
The next approach was using a Segment Tree. We also used the fact that the number of divisors of a positive integer is odd if only the integer is a perfect square. It was taking O(q* logn) time to solve.
The most optimal solution was using the prefix sum of the array elements. The time complexity of this solution was O(n).

Check out this problem - Two Sum Problem

Use Coding Ninjas Studio to practice various DSA questions asked in many interviews. It will assist you in mastering efficient coding techniques, and you will also get interview experiences from people working in big companies.