Open In App

Range Sum Queries and Update with Square Root

Improve
Improve
Like Article
Like
Save
Share
Report

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


Java




import java.util.*;
import java.util.stream.*;
 
class Program {
// Maximum size of input array
static int MAX = 100;
static int[] BIT = new int[MAX + 1];
// structure for queries with members type,
// leftIndex, rightIndex of the query
static class queries { public int type, l, r;
    public queries(int t_val, int l_val, int r_val)
    {
        type = t_val;
        l = l_val;
        r = r_val;
    }
     
}
 
// function for updating the value
static void update(int x, int val, int n) {
    for (; x <= n; x += x & -x) {
        BIT[x] += val;
    }
}
 
// function for calculating the required
// sum between two indexes
static int sum(int x) {
    int s = 0;
    for (; x > 0; x -= x & -x) {
        s += BIT[x];
    }
    return s;
}
 
// function to return answer to queries
static void answerQueries(int[] arr, queries[] que, int n, int q) {
    // Declaring a Set
    HashSet<Integer> s = new HashSet<Integer>();
    for (int i = 1; i < n; i++) {
 
        // inserting indexes of those numbers
        // which are greater than 1
        if (arr[i] > 1)
            s.add(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
                 
                final var val = que[i].l;
                Stream<Integer> st = s.stream().filter(x -> x >= val);
                int it = st.findFirst().orElse(0);
 
                // if it crosses the right index of
                // query or end of set, then break
                if (it == 0 || it > que[i].r)
                    break;
 
                que[i].l = it;
 
                // update the value of arr[i] to
                // its square root
                update(it, (int)Math.sqrt(arr[it]) - arr[it], n);
 
                arr[it] = (int)Math.sqrt(arr[it]);
 
                // if updated value becomes equal to 1
                // remove it from the set
                if (arr[it] == 1)
                    s.remove(it);
 
                // increment the index
                que[i].l++;
            }
        }
 
        // sum query
        else {
            System.out.println(sum(que[i].r) - sum(que[i].l - 1));
        }
    }
}
 
// Driver Code
public static void main(String[] args) {
    int q = 4;
 
    // input array using 1-based indexing
    int[] arr = { 0, 4, 5, 1, 2, 4 };
    int n = arr.length;
 
    // declaring array of structure of type queries
    queries[] que = new queries[q + 1];
    que[0] = new queries(2, 1, 5);
    que[1] = new queries(1, 1, 2);
    que[2] = new queries(1, 2, 4);
    que[3] = new queries(2, 1, 5);
         
 
        // answer the Queries
        answerQueries(arr, que, n, q);
    }
}


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
                # its 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


C#




using System;
using System.Linq;
using System.Collections.Generic;
 
class Program
{
   
    // Maximum size of input array
    const int MAX = 100;
    static int[] BIT = new int[MAX + 1];
 
    // structure for queries with members type,
    // leftIndex, rightIndex of the query
    struct queries { public int type, l, r; }
 
    // function for updating the value
    static void update(int x, int val, int n)
    {
        for (; x <= n; x += x & -x) {
            BIT[x] += val;
        }
    }
 
    // function for calculating the required
    // sum between two indexes
    static int sum(int x)
    {
        int s = 0;
        for (; x > 0; x -= x & -x) {
            s += BIT[x];
        }
        return s;
    }
 
    // function to return answer to queries
    static void answerQueries(int[] arr, queries[] que,
                              int n, int q)
    {
        // Declaring a Set
        var s = new HashSet<int>();
        for (int i = 1; i < n; i++) {
 
            // inserting indexes of those numbers
            // which are greater than 1
            if (arr[i] > 1)
                s.Add(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
                    var it = s.Where(x => x >= que[i].l)
                                 .FirstOrDefault();
 
                    // if it crosses the right index of
                    // query or end of set, then break
                    if (it == 0 || it > que[i].r)
                        break;
 
                    que[i].l = it;
 
                    // update the value of arr[i] to
                    // its square root
                    update(it,
                           (int)Math.Sqrt(arr[it])
                               - arr[it],
                           n);
 
                    arr[it] = (int)Math.Sqrt(arr[it]);
 
                    // if updated value becomes equal to 1
                    // remove it from the set
                    if (arr[it] == 1)
                        s.Remove(it);
 
                    // increment the index
                    que[i].l++;
                }
            }
 
            // sum query
            else {
                Console.WriteLine(sum(que[i].r)
                                  - sum(que[i].l - 1));
            }
        }
    }
 
    // Driver Code
    static void Main(string[] args)
    {
        int q = 4;
 
        // input array using 1-based indexing
        int[] arr = { 0, 4, 5, 1, 2, 4 };
        int n = arr.Length;
 
        // declaring array of structure of type queries
        var que = new queries[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);
    }
}
 
// This code is contributed by phasing17


Javascript




// JavaScript program to calculate sum
// in an interval and update with
// square root
 
// Maximum size of input array
let MAX = 100;
let BIT = new Array(MAX + 1).fill(0);
 
function lower_bound(arr, ele)
{
    for (var i = 0; i < arr.length; i++)
    {
        if (arr[i] >= ele)
            return i;
    }
    return arr.length - 1;
}
 
// structure for queries with members type,
// leftIndex, rightIndex of the query
class queries
{
    constructor(type, l, r)
    {
        this.type = type;
        this.l = l;
        this.r = r;
    }
}
 
// function for updating the value
function update(BIT, x, val, n)
{
    var a = x;
    while (a <= n)
    {
        BIT[a] += val;
        a += (a & -a);
    }
    return BIT;
}
 
// function for calculating the required
// sum between two indexes
function sum(x)
{
    var s = 0;
    var a = x;
    while (a > 0)
    {
        s += BIT[a];
        a -= (a & -a);
    }
 
    return s;
}
 
// function to return answer to queries
function answerQueries(arr, que, n, q)
{
    // Declaring a Set
    let s = new Set();
     
    for (var i = 1; i < n; i++)
    {
        // inserting indexes of those numbers
        // which are greater than 1
        if (arr[i] > 1)
            s.add(i);
        BIT = update(BIT, i, arr[i], n);
    }
    for (var i = 0; i < q; i++)
    {
        // update query
        if (que[i].type == 1)
        {
            while (true)
            {
                var ss = Array.from(s);
                ss.sort();
                 
                // find the left index of query in
                // the set using binary search
                // auto it = s.lower_bound(que[i].l);
                let it = lower_bound(ss, que[i].l);
 
                // if it crosses the right index of
                // query or end of set, then break
                if (it == s.length || ss[it] > que[i].r)
                    break;
                que[i].l = ss[it];
 
                // update the value of arr[i] to
                // its square root
                BIT = update(BIT, ss[it], (Math.pow(arr[ss[it]], 0.5)) - arr[ss[it]], n);
 
                arr[ss[it]] = (Math.pow(arr[ss[it]], 0.5));
 
                // if updated value becomes equal to 1
                // remove it from the set
                if (arr[ss[it]] == 1)
                    arr.splice(ss[it], 1);
 
                // increment the index
                que[i].l += 1;
            }
        }
 
        // sum query
        else
            console.log(Math.floor(sum(que[i].r) - sum(que[i].l - 1)));
    }
}
 
 
// Driver Code
let q = 4;
 
// input array using 1-based indexing
let arr = [0, 4, 5, 1, 2, 4];
let n = arr.length;
// declaring array of structure of type queries
let que = [new queries(2, 1, 5), new queries(1, 1, 2), new queries(1, 2, 4),
new queries(2, 1, 5)];
 
// answer the Queries
answerQueries(arr, que, n, q);
 
// This code is contributed by phasing17


Output

16
9

Complexity Analysis:

  • Time Complexity: O(logN) per query 
  • Auxiliary Space: O(N) 


Last Updated : 30 Jan, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads