Skip to content
Related Articles

Related Articles

Improve Article
Range Sum Queries and Update with Square Root
  • Difficulty Level : Medium
  • Last Updated : 07 Jan, 2021

Given an array A of N integers and number of queries Q. You have to answer two types of queries. 
 

  • Update [l, r] – for every i in range from l to r update Ai with sqrt(Ai), where sqrt(Ai) represents the square root of Ai in integral form.
  • Query [l, r] – calculate the sum of all numbers ranging between l and r in array A.

Prerequisite :Binary Indexed Trees | Segment Trees
Examples: 
 

Input: A[] = { 4, 5, 1, 2, 4 }, Q = {{2, 1, 5}, {1, 1, 2}, {1, 2, 4}, {2, 1, 5}} 
Output: 16 

Considering 1-based indexing, first query is to calculate sum of numbers from A1 to A5 
which is 4 + 5 + 1 + 2 + 4 = 16. 
Second query is to update A1 to A2 with its square root. Now, array becomes A[] = { 2, 2, 1, 2, 4 }. 
Similarly, third query is to update A2 to A4 with its square root. Now, array becomes A[] = { 2, 1, 1, 1, 4 }. 
Fourth query is to calculate sum of numbers from A1 to A5 which is 2 + 1 + 1 + 1 + 4 = 9.
Input: A[] = { 4, 9, 25, 36 }, Q = {{1, 2, 4}, {2, 1, 4}} 
Output: 18 
 

 

Naive Approach: A simple solution is to run a loop from l to r and calculate sum of elements in the given range. To update a value, simple replace arr[i] with its square root, i.e., arr[i] = sqrt[arr[i]].
Efficient Approach: The idea is to reduce the time complexity for each query and update operation to O(logN). Use Binary Indexed Trees (BIT) or Segment Trees. Construct a BIT[] array and have two functions for query and update operation. Now, for each update operation the key observation is that the number 1 will have 1 as its square root, so if it exists in the range of update query, it doesn’t need to be updated. We will use a set to store the index of only those numbers which are greater than 1 and use binary search to find the l index of the update query and increment the l index until every element is updated in range of that update query. If the arr[i] has 1 as its square root then after updating it, remove it from the set as it will always be 1 even after any next update query. For sum query operation, simply do query(r) – query(l – 1).
Below is the implementation of the above approach:
 



CPP




// CPP program to calculate sum
// in an interval and update with
// square root
#include <bits/stdc++.h>
using namespace std;
 
// Maximum size of input array
const int MAX = 100;
 
int BIT[MAX + 1];
 
// structure for queries with members type,
// leftIndex, rightIndex of the query
struct queries {
    int type, l, r;
};
 
// function for updating the value
void update(int x, int val, int n)
{
    for (x; x <= n; x += x & -x) {
        BIT[x] += val;
    }
}
 
// function for calculating the required
// sum between two indexes
int sum(int x)
{
    int s = 0;
    for (x; x > 0; x -= x & -x) {
        s += BIT[x];
    }
    return s;
}
 
// function to return answer to queries
void answerQueries(int arr[], queries que[], int n, int q)
{
    // Declaring a Set
    set<int> s;
    for (int i = 1; i < n; i++) {
 
        // inserting indexes of those numbers
        // which are greater than 1
        if (arr[i] > 1)
            s.insert(i);
        update(i, arr[i], n);
    }
 
    for (int i = 0; i < q; i++) {
 
        // update query
        if (que[i].type == 1) {
            while (true) {
 
                // find the left index of query in
                // the set using binary search
                auto it = s.lower_bound(que[i].l);
 
                // if it crosses the right index of
                // query or end of set, then break
                if (it == s.end() || *it > que[i].r)
                    break;
 
                que[i].l = *it;
 
                // update the value of arr[i] to
                // its square root
                update(*it, (int)sqrt(arr[*it]) - arr[*it], n);
 
                arr[*it] = (int)sqrt(arr[*it]);
 
                // if updated value becomes equal to 1
                // remove it from the set
                if (arr[*it] == 1)
                    s.erase(*it);
 
                // increment the index
                que[i].l++;
            }
        }
 
        // sum query
        else {
            cout << (sum(que[i].r) - sum(que[i].l - 1)) << endl;
        }
    }
}
 
// Driver Code
int main()
{
    int q = 4;
 
    // input array using 1-based indexing
    int arr[] = { 0, 4, 5, 1, 2, 4 };
    int n = sizeof(arr) / sizeof(arr[0]);
 
    // declaring array of structure of type queries
    queries que[q + 1];
 
    que[0].type = 2, que[0].l = 1, que[0].r = 5;
    que[1].type = 1, que[1].l = 1, que[1].r = 2;
    que[2].type = 1, que[2].l = 2, que[2].r = 4;
    que[3].type = 2, que[3].l = 1, que[3].r = 5;
 
    // answer the Queries
    answerQueries(arr, que, n, q);
 
    return 0;
}

Python3




# Python program to calculate sum
# in an interval and update with
# square root
from typing import List
import bisect
from math import sqrt, floor
 
# Maximum size of input array
MAX = 100
BIT = [0 for _ in range(MAX + 1)]
 
# structure for queries with members type,
# leftIndex, rightIndex of the query
class queries:
    def __init__(self, type: int = 0, l: int = 0, r: int = 0) -> None:
        self.type = type
        self.l = l
        self.r = r
 
# function for updating the value
def update(x: int, val: int, n: int) -> None:
    a = x
    while a <= n:
        BIT[a] += val
        a += a & -a
 
# function for calculating the required
# sum between two indexes
def sum(x: int) -> int:
    s = 0
    a = x
    while a:
        s += BIT[a]
        a -= a & -a
 
    return s
 
# function to return answer to queries
def answerQueries(arr: List[int], que: List[queries], n: int, q: int) -> None:
 
    # Declaring a Set
    s = set()
    for i in range(1, n):
 
        # inserting indexes of those numbers
        # which are greater than 1
        if (arr[i] > 1):
            s.add(i)
        update(i, arr[i], n)
 
    for i in range(q):
 
        # update query
        if (que[i].type == 1):
            while True:
                ss = list(sorted(s))
                 
                # find the left index of query in
                # the set using binary search
                # auto it = s.lower_bound(que[i].l);
                it = bisect.bisect_left(ss, que[i].l)
 
                # if it crosses the right index of
                # query or end of set, then break
                if it == len(s) or ss[it] > que[i].r:
                    break
                que[i].l = ss[it]
 
                # update the value of arr[i] to
                # itss square root
                update(ss[it], floor(sqrt(arr[ss[it]]) - arr[ss[it]]), n)
 
                arr[ss[it]] = floor(sqrt(arr[ss[it]]))
 
                # if updated value becomes equal to 1
                # remove it from the set
                if (arr[ss[it]] == 1):
                    s.remove(ss[it])
 
                # increment the index
                que[i].l += 1
 
        # sum query
        else:
            print(sum(que[i].r) - sum(que[i].l - 1))
 
 
# Driver Code
if __name__ == "__main__":
 
    q = 4
 
    # input array using 1-based indexing
    arr = [0, 4, 5, 1, 2, 4]
    n = len(arr)
    # declaring array of structure of type queries
    que = [queries() for _ in range(q + 1)]
 
    que[0].type, que[0].l, que[0].r = 2, 1, 5
    que[1].type, que[1].l, que[1].r = 1, 1, 2
    que[2].type, que[2].l, que[2].r = 1, 2, 4
    que[3].type, que[3].l, que[3].r = 2, 1, 5
 
    # answer the Queries
    answerQueries(arr, que, n, q)
 
# This code is contributed by sanjeev2552
Output: 
16
9

 

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with industry experts, please refer Geeks Classes Live




My Personal Notes arrow_drop_up
Recommended Articles
Page :