Powerful Array

 Problem Link

Algorithms Used: square root decomposition, sorting

In this problem, we need to answer some queries of the form (l r) by reporting the value of val * freq² where val is the value of a number in the subarray[l...r] and freq is the number of times the value 'val' shows up in the subarray.

A simple, straightforward way of doing this is to iterate over each element of the subarray, update a map storing the frequency of a value and then calculate the answer using these values and frequencies. This algorithm has a time complexity of O(NQlog(N)), because it take O(NlogN) time to iterate over and insert each number into the map from each subarray and there are Q queries. This will easily TLE.

Instead, we can sort the queries. Intuitively, we might sort the queries by their left values and then use the two pointer method: shift the left pointer to the right whenever the subarray left values increase, and shift the right pointer either left or right depending on the scope of the previous query and the current query. However, even this will TLE because in the worst case, the right pointer will have to move N-1 places for every query (the right value for the first query could be 1, then for the second it could be N, then for the third it could be 1 and so on. The right pointer would keep switching between 1 and N and each switch takes O(N) time) which is repeated Q times yielding a time complexity of O(NQ)). 

What we need is a way to order the queries such that the right pointer shifts a smaller distance for each query. We can do this by sorting the queries in ascending order of their left values, and sort the queries in each block of √Q queries in ascending order of their right values. This way, the left pointer moves a maximum of N times and the right pointer moves a maximum of 2N√Q times because the right pointer moves a maximum N times forward when processing the queries in a block and N times backward to start answering the queries in the next block.

Code

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
using namespace std;
using ll = long long;
const int maxa = 1e6 + 1;
ll freq[maxa];
ll ans = 0;
void add_num(ll x) {
    ans -= (freq[x] * freq[x] * x);
    freq[x]++;
    ans += (freq[x] * freq[x] * x);
}
void rem_num(int x) {
    ans -= (freq[x] * freq[x] * x);
    freq[x]--;
    ans += (freq[x] * freq[x] * x);
}
int sz;
bool comp (pair <pair <ll, ll>, ll> p1, pair <pair <ll, ll>, ll> p2) {
    if (p1.first.first / sz == p2.first.first / sz)
        return p1.first.second < p2.first.second;
    return p1.first.first / sz < p2.first.first / sz;
}
int main() {
    int n, q;
    cin >> n >> q;
    vector <ll> a(n+1);
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    sz = sqrt(n);
    vector <pair <pair <ll, ll>, int>> queries(q);
    for (int i = 0; i < q; i++) {
        cin >> queries[i].first.first >> queries[i].first.second;
        queries[i].second = i;
    }
    sort(queries.begin(), queries.end(), comp);
    int prevl = 1, prevr = 0;
    vector <ll> p_ans(q);
    for (int i = 0; i < q; i++) {
        int l = queries[i].first.first, r = queries[i].first.second;
        int id = queries[i].second;
       
        while (prevr < r) {
            add_num(a[++prevr]);
        }
        while (prevr > r) {
            rem_num(a[prevr--]);
        }
        while (prevl < l) {
            rem_num(a[prevl++]);
        }
        while (prevl > l) {
            add_num(a[--prevl]);
        }
        p_ans[id] = ans;
    }
    for (int i = 0; i < q; i++)
        cout << p_ans[i] << '\n';
}

Comments