EQSUM - Editorial

editorial
, , , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Cakewalk

PREREQUISITES:

Maps/sorting

PROBLEM:

Given an array A, count the number of pairs such that A_i - A_j = i - j

EXPLANATION:

To solve this problem, we rearrange the equation A_i - A_j = i - j a bit to obtain the equation

A_i - i = A_j - j

instead.

So, if we define B_i = A_i - i, we’re now looking for the number of pairs (i, j) such that i \lt j and B_i = B_j.

Counting the number of equal pairs of elements in an array of length N can be done in \mathcal{O}(N\log N) in several ways.
Perhaps the easiest way is to use a map (dictionary in Python.)

Let M be a map, initially empty.
For each i = 1, 2, 3, \ldots, N in order,

  • First, add M[B_i] to the answer.
    This will count the number of valid pairs whose second index equals i.
  • Then, add 1 to M[B_i].
    This is so that index i will be added for future occurrences of the same value.

In the end, print the answer.

Standard map operations are \mathcal{O}(\log N) in C++, so this is an \mathcal{O}(N\log N) solution which is perfectly fine.

There are other solutions as well: for example you can use sorting and apply this algorithm.

TIME COMPLEXITY:

\mathcal{O}(N\log N) per testcase.

CODE:

Editorialist's code (C++) 
// #include <bits/allocator.h>
// #pragma GCC optimize("Ofast,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());

int main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    int t; cin >> t;
    while (t--) {
        int n; cin >> n;
        vector a(n, 0);
        for (int &x : a) cin >> x;

        map<int, int> mp;
        int ans = 0;
        for (int i = 0; i < n; ++i) {
            ans += mp[a[i] - i];
            ++mp[a[i] - i];
        }
        cout << ans << '\n';
    }
}