Skip to content
Related Articles
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 ``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);` `    ``// declaring array of structure of type queries``    ``queries que[q + 1];` `    ``que.type = 2, que.l = 1, que.r = 5;``    ``que.type = 1, que.l = 1, que.r = 2;``    ``que.type = 1, que.l = 2, que.r = 4;``    ``que.type = 2, que.l = 1, que.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