Range Sum Queries and Update with Square Root


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
9
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:

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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;
}

chevron_right


Output:

16
9


My Personal Notes arrow_drop_up