PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: Satyam
Testers: Nishank Suresh, Tejas Pandey
Editorialist: Nishank Suresh
DIFFICULTY:
3042
PREREQUISITES:
Basic algebraic manipulation, frequency arrays
PROBLEM:
Given a permutation P of \{1, 2, \ldots, N\}, count the number of pairs (i, j) such that 1 \leq i \leq j \leq N and i+j divides P_i + P_j
EXPLANATION:
The condition “i+j divides P_i + P_j” can be rewritten as “there exists an integer k such that P_i + P_j = k\cdot (i+j)”.
Notice that the above equation can be rewritten as P_i - k\cdot i = k\cdot j - P_j = - (P_j - k\cdot j).
In particular, if the value of k is fixed, it’s not hard to find, in linear time, all pairs (i, j) that achieve divisibility using this k.
How?
Iterate over j from 1 to N. Also maintain a frequency array freq, which is initially empty.
When at an index j,
- Add 1 to freq[P_j - k\cdot j]
- The number of i \leq j that satisfy the condition is then exactly freq[k\cdot j - P_j]
This can be done in \mathcal{O}(N) time.
However, repeating the above algorithm for each k is too slow, since k can be anything at all and each run takes \mathcal{O}(N) time.
Let’s make a few observations to reduce the work we do.
The first thing to notice is that P is a permutation, which means that P_i \leq N for every i.
In particular, P_i + P_j \leq 2N. Combine this with the fact that i+j \geq 2, and you can see that any valid k must satisfy 1 \leq k \leq N.
We now have a solution in \mathcal{O}(N^2), which is still too slow.
Now, suppose we fix the value of k. Do we really need to look at every element of the array?
No!
Again, let’s look at P_i + P_j = k\cdot (i+j). We know P_i + P_j \leq 2N, so k\cdot(i+j) \leq 2N.
This tells us that i+j \leq \frac{2N}{k}: in particular, j \lt \frac{2N}{k}.
So, we don’t need to iterate across the whole array: once k is fixed, only run the initial algorithm on the first \frac{2N}{k} elements.
In fact, this optimization immediately improves our complexity: the number of operations we perform is
which is \mathcal{O}(N\log N), being a harmonic summation.
Note that the constraints in this problem are somewhat high, and so using a map/unordered_map/dict to maintain your frequency array will most likely result in a TLE verdict because of the high constant factor and/or an extra \log factor.
Instead, you can maintain a global frequency array of size 4N and use an offset to account for negative numbers.
That is, if you want to add 1 to the frequency of x, you instead add 1 to freq[x + 2N].
This shifts the indices of the frequency table from [-2N, 2N] to [0, 4N].
Once you’ve used the frequency table for a specific k, make sure to reset it properly.
If working with an offset-shifted array like this is new to you, take a look at the code linked below.
TIME COMPLEXITY:
\mathcal{O}(N \log N) per testcase.
CODE:
Setter's code (C++)
#include <bits/stdc++.h>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
using namespace std;
#define ll long long
const ll INF_MUL=1e13;
const ll INF_ADD=1e18;
#define pb push_back
#define mp make_pair
#define nline "\n"
#define f first
#define s second
#define pll pair<ll,ll>
#define all(x) x.begin(),x.end()
#define vl vector<ll>
#define vvl vector<vector<ll>>
#define vvvl vector<vector<vector<ll>>>
#ifndef ONLINE_JUDGE
#define debug(x) cerr<<#x<<" "; _print(x); cerr<<nline;
#else
#define debug(x);
#endif
void _print(ll x){cerr<<x;}
void _print(char x){cerr<<x;}
void _print(string x){cerr<<x;}
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
template<class T,class V> void _print(pair<T,V> p) {cerr<<"{"; _print(p.first);cerr<<","; _print(p.second);cerr<<"}";}
template<class T>void _print(vector<T> v) {cerr<<" [ "; for (T i:v){_print(i);cerr<<" ";}cerr<<"]";}
template<class T>void _print(set<T> v) {cerr<<" [ "; for (T i:v){_print(i); cerr<<" ";}cerr<<"]";}
template<class T>void _print(multiset<T> v) {cerr<< " [ "; for (T i:v){_print(i);cerr<<" ";}cerr<<"]";}
template<class T,class V>void _print(map<T, V> v) {cerr<<" [ "; for(auto i:v) {_print(i);cerr<<" ";} cerr<<"]";}
typedef tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;
typedef tree<ll, null_type, less_equal<ll>, rb_tree_tag, tree_order_statistics_node_update> ordered_multiset;
typedef tree<pair<ll,ll>, null_type, less<pair<ll,ll>>, rb_tree_tag, tree_order_statistics_node_update> ordered_pset;
//--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
const ll MOD=998244353;
const ll MAX=5000500;
vector<ll> freqpos(MAX,0),freqneg(MAX,0);
void solve(){
ll n; cin>>n;
vector<ll> p(n+5,0);
ll ans=0;
for(ll i=1;i<=n;i++){
cin>>p[i];
if(p[i]%i){
;
}
else{
ans++;
}
}
for(ll k=1;k<=2*n;k++){
ll till=min(n,(2*n)/k);
ll zero=0;
for(ll i=1;i<=till;i++){
ll now=p[i]-i*k;
assert(abs(now)<MAX);
if(now==0){
ans+=zero++;
}
else if(now>0){
ans+=freqneg[now];
freqpos[now]++;
}
else{
ans+=freqpos[-now];
freqneg[-now]++;
}
}
for(ll i=1;i<=till;i++){
ll now=p[i]-i*k;
if(now==0){
;
}
else if(now>0){
freqpos[now]--;
}
else{
freqneg[-now]--;
}
}
}
cout<<ans<<nline;
return;
}
int main()
{
ios_base::sync_with_stdio(false);
cin.tie(NULL);
#ifndef ONLINE_JUDGE
freopen("input.txt", "r", stdin);
freopen("output.txt", "w", stdout);
freopen("error.txt", "w", stderr);
#endif
ll test_cases=1;
cin>>test_cases;
while(test_cases--){
solve();
}
cout<<fixed<<setprecision(10);
cerr<<"Time:"<<1000*((double)clock())/(double)CLOCKS_PER_SEC<<"ms\n";
}
Editorialist's code (Python)
import sys
input = sys.stdin.readline
for _ in range(int(input())):
n = int(input())
p = list(map(int, input().split()))
ans = 0
freq = [0 for _ in range(4*n + 10)]
offset = 2*n
for k in range(1, n+1):
lim = min(2*n // k, n) + 1
for j in range(1, lim):
freq[p[j-1] - k*j + offset] += 1
ans += freq[k*j - p[j-1] + offset]
for j in range(1, lim):
freq[p[j-1] - k*j + offset] -= 1
print(ans)