Skip to content
Related Articles

Related Articles

Queries to calculate sum of squares of array elements over range of indices [L, R] with updates
  • Last Updated : 19 Apr, 2021

Given an array arr[] consisting of N integers, an integer X, and a 2D array queries[][] consisting of queries of the following two types:

  • (L, R, 0): Increment all array elements over the range [L, R] by a value X.
  • (L, R, 1): Set all array elements to X over the range [L, R].
  • (L, R, 2): Print the sum of the square of the elements over the range [L, R].

Examples:

Input: arr[] = {1, 2, 3, 4}, X = 2, queries[][] = {{1, 3, 0}, {1, 2, 2}, {3, 3, 1}, {2, 3, 2}}
Output: 41 29
Explanation:
Query 1: Increment all array elements present in the range of indices [1, 3] by X ( = 2). The array modifies to {1, 4, 5, 6, 5}.
Query 2: Sum of squares of array elements present in the range of indices [1, 2] = 4*4 + 5*5 = 41.
Query 3: Updating all array elements to X(= 2) in the range of indices [3, 3]. The array modifies to {1, 4, 5, 2, 5}.
Query 4: Sum of squares of array elements in the range of indices [2, 3] = 5 * 5 + 2 * 2 = 29.

Input: arr[] = {2, 4, 1, 5}, X = 3, queries[][] = {{1, 3, 1}, {1, 3, 0}, {2, 3, 0}}
Output: 27 18

Naive Approach: The simplest approach is to traverse the given array over the given ranges for each query and perform the corresponding queries on the given array and print the result accordingly.



Below is the implementation of the above approach:

C++




// C++ program for the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Function to find the sum of squares
// of elements present in the range of
// indices [l, r] in the array v[]
long long rangesum(vector<int>& v,
                   int l, int r)
{
    long long sum = 0;
 
    // Iterate over range l to r
    for (int i = l; i <= r; i++) {
 
        // Add square of current
        // element to the sum
        sum += v[i] * v[i];
    }
 
    // Return the sum
    return sum;
}
 
// Function to perform given queries
void rangeQuery(vector<int>& v, int l,
                int r, int update,
                int type)
{
    // For type 1 update
    if (type == 1) {
 
        // Iterate over range l to r
        // and update the array elements
        for (int i = l; i <= r; i++) {
            v[i] = update;
        }
    }
 
    // For type 0 update
    else {
 
        // Iterate over range l to r
        // and update the elements
        for (int i = l; i <= r; i++) {
            v[i] += update;
        }
    }
}
 
// Function to print the result
// for the given queries
void printAnswer(vector<vector<int> > query,
                 vector<int> a, int X)
{
 
    // Iterate over the query[]
    for (int i = 0;
         i < query.size(); i++) {
 
        int l = query[i][0];
        int r = query[i][1];
        if (query[i][2] == 0) {
 
            // Add by X modified array
            // will be considered for
            // the operation
            rangeQuery(a, l, r, X, 0);
        }
        else if (query[i][2] == 1) {
 
            // Substitute by X
            rangeQuery(a, l, r, X, 1);
        }
        else {
 
            // Print the range sum
            cout << rangesum(a, l, r)
                 << " ";
        }
    }
}
 
// Driver Code
int main()
{
    vector<int> arr = { 1, 2, 3, 4, 5 };
    int X = 2;
    vector<vector<int> > query = {
        { 1, 3, 0 }, { 1, 2, 2 },
        { 3, 4, 1 }, { 2, 3, 2 }
    };
 
    printAnswer(query, arr, X);
 
    return 0;
}

Java




// Java program for the above approach
import java.util.*;
 
class GFG{
     
// Function to find the sum of squares
// of elements present in the range of
// indices [l, r] in the array v[]
static long rangesum(List<Integer> v, int l, int r)
{
    long sum = 0;
 
    // Iterate over range l to r
    for(int i = l; i <= r; i++)
    {
         
        // Add square of current
        // element to the sum
        sum += v.get(i) * v.get(i);
    }
 
    // Return the sum
    return sum;
}
 
// Function to perform given queries
static void rangeQuery(List<Integer> v, int l, int r,
                       int update, int type)
{
     
    // For type 1 update
    if (type == 1)
    {
 
        // Iterate over range l to r
        // and update the array elements
        for(int i = l; i <= r; i++)
        {
            v.set(i, update);
        }
    }
 
    // For type 0 update
    else
    {
 
        // Iterate over range l to r
        // and update the elements
        for(int i = l; i <= r; i++)
        {
            v.set(i, v.get(i) + update);
        }
    }
}
 
// Function to print the result
// for the given queries
static void printAnswer(int[][] query, List<Integer> a,
                        int X)
{
 
    // Iterate over the query[]
    for(int i = 0; i < query.length; i++)
    {
        int l = query[i][0];
        int r = query[i][1];
         
        if (query[i][2] == 0)
        {
             
            // Add by X modified array
            // will be considered for
            // the operation
            rangeQuery(a, l, r, X, 0);
        }
        else if (query[i][2] == 1)
        {
             
            // Substitute by X
            rangeQuery(a, l, r, X, 1);
        }
        else
        {
             
            // Print the range sum
            System.out.print(rangesum(a, l, r) + " ");
        }
    }
}
 
// Driver code
public static void main(String[] args)
{
    List<Integer> arr = new ArrayList<>(
        Arrays.asList(1, 2, 3, 4, 5));
    int X = 2;
    int[][] query = { { 1, 3, 0 },
                      { 1, 2, 2 },
                      { 3, 4, 1 },
                      { 2, 3, 2 } };
 
    printAnswer(query, arr, X);
}
}
 
// This code is contributed by offbeat

C#




// C# program for the above approach
using System;
using System.Collections.Generic;
 
class GFG{
 
// Function to find the sum of squares
// of elements present in the range of
// indices [l, r] in the array v[]
static long rangesum(List<int> v, int l, int r)
{
    long sum = 0;
 
    // Iterate over range l to r
    for(int i = l; i <= r; i++)
    {
         
        // Add square of current
        // element to the sum
        sum += v[i] * v[i];
    }
 
    // Return the sum
    return sum;
}
 
// Function to perform given queries
static void rangeQuery(List<int> v, int l, int r,
                       int update, int type)
{
 
    // For type 1 update
    if (type == 1)
    {
 
        // Iterate over range l to r
        // and update the array elements
        for(int i = l; i <= r; i++)
        {
            v[i] = update;
        }
    }
 
    // For type 0 update
    else
    {
 
        // Iterate over range l to r
        // and update the elements
        for(int i = l; i <= r; i++)
        {
            v[i] = v[i] + update;
        }
    }
}
 
// Function to print the result
// for the given queries
static void printAnswer(int[,] query, List<int> a,
                        int X)
{
 
    // Iterate over the query[]
    for(int i = 0; i < query.GetLength(0); i++)
    {
        int l = query[i, 0];
        int r = query[i, 1];
 
        if (query[i, 2] == 0)
        {
 
            // Add by X modified array
            // will be considered for
            // the operation
            rangeQuery(a, l, r, X, 0);
        }
        else if (query[i, 2] == 1)
        {
 
            // Substitute by X
            rangeQuery(a, l, r, X, 1);
        }
        else
        {
 
            // Print the range sum
            Console.Write(rangesum(a, l, r) + " ");
        }
    }
}
 
// Driver code
public static void Main(string[] args)
{
    List<int> arr = new List<int>{ 1, 2, 3, 4, 5 };
    int X = 2;
    int[,] query = { { 1, 3, 0 },
                     { 1, 2, 2 },
                     { 3, 4, 1 },
                     { 2, 3, 2 } };
 
    printAnswer(query, arr, X);
}
}
 
// This code is contributed by ukasp
Output: 
41 29

 

Time Complexity: O(Q*N)
Auxiliary Space: O(N + Q)

Efficient Approach: The above approach can be optimized using Segment Tree. Follow the steps below to solve the problem:

(b + K)2 + (c + K)2 + (d + K)2 = (b2 + c2 + d2) + 3 * K2 + 2 * K * (b + c + d), 
which can be generalized as the sum of squares + (high – low + 1) * X * X + sum of elements.

  • On substituting X = K in this range, (K)2 + (K)2 + (K)2 = 3 * K2, which can be generalized to: (high – low + 1) * X * X.
  • Update the sum of the squares of array elements over the range [L, R] using the above formula and print the range sum accordingly for each query of type 2.

Below is the implementation of the above approach:

C++




// C++ program of the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
// Used to create Segment Tree
vector<long long> seg(400000, 0),
    dseg(400000, 0),
    a(100000);
 
// Stores the segment sum
vector<long long> segSum(400000, 0);
 
// Building the segment tree
void build(int ind, int low, int high)
{
    // If low is the same as high
    if (low == high) {
 
        // Update the segment node
        seg[ind] = a[low] * a[low];
        segSum[ind] = a[low];
        return;
    }
 
    // Find the mid
    int mid = (low + high) / 2;
 
    // Recursively call for
    // the two halves of the tree
    build(2 * ind + 1, low, mid);
    build(2 * ind + 2, mid + 1, high);
 
    // Update the segment node
    seg[ind] = seg[2 * ind + 1]
               + seg[2 * ind + 2];
 
    // Update the segment sum
    segSum[ind] = segSum[2 * ind + 1]
                  + segSum[2 * ind + 2];
}
 
// Function to find the sum of squares
// of array elements in the given range
long long rangesum(int ind, int low,
                   int high, int l,
                   int r)
{
    // Pending updates if any
    if (dseg[ind] != 0) {
 
        // Update the square
        // over segment node
        seg[ind] += (high - low + 1)
                        * dseg[ind]
                        * dseg[ind]
                    + (2 * dseg[ind]
                       * segSum[ind]);
 
        // Update the square over
        // the segment sum
        segSum[ind] += (high - low + 1)
                       * dseg[ind];
 
        // If low and high are different
        if (low != high) {
            dseg[2 * ind + 1] += dseg[ind];
            dseg[2 * ind + 2] += dseg[ind];
        }
        dseg[ind] = 0;
    }
 
    // Out of bounds
    if (r < low || l > high
        || low > high) {
        return 0;
    }
 
    // Completely within the range
    if (l <= low && r >= high) {
        return seg[ind];
    }
 
    // Partially overlapping
    int mid = (low + high) / 2;
 
    // Return the range sum
    return rangesum(2 * ind + 1,
                    low, mid, l, r)
           + rangesum(2 * ind + 2,
                      mid + 1, high,
                      l, r);
}
 
// Function to perform the given queries
void rangeQuery(int ind, int low,
                int high, int l, int r,
                int update, int type)
{
    // Pending updates if any
    if (dseg[ind] != 0) {
 
        // Update the segment node
        seg[ind]
            += (high - low + 1)
               * dseg[ind] * dseg[ind];
 
        // Update the segment sum
        segSum[ind] += (high - low + 1)
                       * dseg[ind];
 
        if (low != high) {
            dseg[2 * ind + 1] += dseg[ind];
            dseg[2 * ind + 2] += dseg[ind];
        }
 
        dseg[ind] = 0;
    }
 
    // Out of bounds
    if (r < low || l > high
        || low > high) {
        return;
    }
 
    // Completely within the range
    if (l <= low && r >= high) {
 
        // Substitute with X
        if (type == 1) {
 
            // Updating the current
            // index for sum of square
            seg[ind] = (high - low + 1)
                       * update * update;
 
            // Updating the current
            // index for sum
            segSum[ind] = (high - low + 1)
                          * update;
        }
        else {
 
            // Add X updating the current
            // index for sum of square
            seg[ind] += (high - low + 1)
                            * update * update
                        + 2 * update
                              * segSum[ind];
 
            // Updating the current
            // index for sum
            segSum[ind] += (high - low + 1)
                           * update;
        }
 
        // Lazy update
        if (low != high) {
            dseg[2 * ind + 1] += update;
            dseg[2 * ind + 2] += update;
        }
        return;
    }
 
    // Partially overlapping
    int mid = (low + high) / 2;
    rangeQuery(2 * ind + 1, low,
               mid, l, r, update,
               type);
 
    rangeQuery(2 * ind + 2, mid + 1,
               high, l, r, update,
               type);
 
    // Updating sum of squares
    // and sum while bactracking
    seg[ind] = seg[2 * ind + 1]
               + seg[2 * ind + 2];
 
    segSum[ind] = segSum[2 * ind + 1]
                  + segSum[2 * ind + 2];
}
 
// Function to print the answer
// for the given queries
void printAnswer(
    vector<vector<int> > query,
    int N, int X)
{
    // Build Segment tree
    build(0, 0, N - 1);
 
    // Traverse queries
    for (int i = 0;
         i < query.size(); i++) {
 
        int l = query[i][0];
        int r = query[i][1];
        if (query[i][2] == 0) {
 
            // Add by X, modified
            // array will be considered
            // for operation
            rangeQuery(0, 0, N - 1, l, r, X, 0);
        }
        else if (query[i][2] == 1) {
 
            // Substitute by X modified
            // array will be considered
            // for operation
            rangeQuery(0, 0, N - 1,
                       l, r, X, 1);
        }
        else {
            cout << rangesum(0, 0, N - 1,
                             l, r)
                 << " ";
        }
    }
}
 
// Driver Code
int main()
{
    a = { 1, 2, 3, 4 };
    int X = 2;
    int N = (int)a.size();
 
    // Given queries
    vector<vector<int> > query = {
        { 1, 3, 0 }, { 1, 2, 2 },
        { 3, 3, 1 }, { 2, 3, 2 }
    };
 
    printAnswer(query, N, X);
 
    return 0;
}
Output: 
41 29

 

Time Complexity: O(Q*log N)
Auxiliary Space: O(N)

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 and Geeks Classes Live USA

My Personal Notes arrow_drop_up
Recommended Articles
Page :