PROBLEM LINK
Setter & Editorialist: Mayuresh Patle
Testers: Uttkarsh Bawankar, Vallabh Deshpande
DIFFICULTY
EASY
PREREQUISITES
Prefix Sum, Binary Search
PROBLEM
You are given an array A of length N, containing values in range [0,2000], having weight W_i associated with i^{th} position of the array (for 1\leq i \leq N). You have to peform Q queries of two types:
- Type 1: 1 \ X, in this query, you have to add X to current weights of all position.
- Type 2: 2 \ V \ L \ R, in this query, you have to print the sum of weights of all positions in range [L,R] which have value V in array A.
QUICK EXPLANATION
Maintain a variable, say extra to keep track of total increment. For type 1 queries, simply add X to extra. For each i \in [0,2001], maintain 2 arrays, one array should contain the positions of elements having value i and other array should contain the prefix sum of weights corresponding to the saved positions. For each query of type 2, use binary search to get the range of indices of favourable positions in the array maintained for value V, print the difference of prefix sums for new range of indices +\ count*extra, where count = count of indices in the new range.
EXPLANATION
Let’s jump directly to the solution!
Let’s say that pos[i] is an array that contains all the positions in array A which have element value i, i.e. pos[i] contains all positions such that A[position] = i. Let’s say another array pre[i] contains the prefix sums corresponding to value i, such that pre[i][j]= \sum_{x=0}^{x=j}{W[pos[i][x]]}
We have to calculate and store this information for all i such that 0 \leq i \leq 2000. To avoid the overhead of handling exceptions, let’s initialize pos[i][0]=0 and pre[i][0]=0.
Assume that arr.back() gives the value of last element of array arr, and arr.push\_back(x) adds new elelment x at the end of arr. Now, to store the above information, we just have to do the following operations
for i = 1 to N
pos[A[i]].push_back(i)
pre[A[i]].push_back(pre[A[i]].back() + W[i])
Now we are ready with the structure that was discussed above. Next task is to handle the queries.
Consider the qeries of type 2, we have the values of V, L and R, to handle this query, we somehow need to find largest index rt such that pos[V][rt] \leq R and largest index lt such that pos[V][lt] \lt L. This can be done in logarithmic time using binary search since the array pos[i] is already sorted.
Now , our answer is pre[V][rt] - pre[V][lt]
Great! We are done with type 2 query!
Wait, wait, wait!!!
What if some type 1 queries were asked earlier?
We see that the result of Type 2 queries depends previous type 1 queries. Now, how to handle type 1 queries? Should we use Segment Tree or Fenwick Tree?
Don’t worry, we don’t need Segment Tree or Fenwick Tree here.
Since all weights are being updated, we can simply use a variable, say extra, to keep track of the total increment.
Initially, extra=0. If we encounter any query of type 1, we can simply do
extra = extra + X
Now, final task is to calculate the correct answer for Type 2 queries. We observe that we just need to add extra*count to the previously calculated answer, where count = number of elements to be considered.
Great we have the exact answer now, i.e.
pre[V][rt] - pre[V][lt] + extra*(rt - lt)
For each testcase:
Time Complexity = O(Q*logN)
Auxiliary Space = O(N)
SOLUTION
Setter's Solution
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector <ll> vll;
#define pb push_back
#define all(vec) vec.begin(),vec.end()
int main()
{
ios_base::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
ll T;
cin>>T;
while(T--)
{
vector <vll> pos(2001,vll(1,0)), pre(2001,vll(1,0));
ll extra=0,n,q,i,type;
cin>>n>>q;
ll a[n],w[n];
for(ll& i:a) cin>>i; //Input array A
for(ll& i:w) cin>>i; //Input weights
for(i=0;i<n;++i)
{
pos[a[i]].pb(i+1); //Add current index i in the positions of a[i]
pre[a[i]].pb(pre[a[i]].back()+w[i]); //Store the prefix sum of weights till current index
}
while(q--)
{
cin>>type;
//Type 1 Queries
if(type==1)
{
ll x;
cin>>x;
extra+=x;
continue;
}
//Type 2 Queries
ll v,l,r;
cin>>v>>l>>r;
auto lt=lower_bound(all(pos[v]),l);
auto rt=upper_bound(all(pos[v]),r);
--lt,--rt;
l=lt-pos[v].begin();
r=rt-pos[v].begin();
cout<<pre[v][r]-pre[v][l]+(r-l)*extra<<"\n";
}
}
return 0;
}