PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: ro27
Tester: jay_1048576
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
None
PROBLEM:
Given an array A, find the number of pairs of indices i \lt j such that:
- The sequence (A_i-A_j, A_i+A_j, A_i\times A_j) is an arithmetic progression.
EXPLANATION:
Suppose, for some indices i and j, the sequence (A_i-A_j, A_i+A_j, A_i\times A_j) is an arithmetic progression.
Let’s see what that tells us about the values of A_i and A_j.
A bit of algebra
Three numbers being in AP is equivalent to saying that twice the middle number equals the sum of the other two.
So, we have 2\cdot (A_i + A_j) = (A_i - A_j) + (A_i\cdot A_j).
Simplifying this a bit,
Now, we use the fact that all the A_i are positive integers.
The above equation tells us that A_i - 3 must be a factor of A_i.
However, this can only happen when A_i = 4, or A_i = 6.
For A_i \geq 7, A_i-3 is too large to be a proper divisor of A_i.
When A_i = 4, we get A_j = 4.
When A_i = 6, we get A_j = 2.
(Technically, A_i = 2 also has 2-3 = -1 as a divisor, but that would result in negative A_j so it can be ignored).
So, the only valid pairs of elements are (6, 2) and (4, 4)!
So, the problem reduces to counting the number of pairs (i, j) such that (A_i, A_j) = (4, 4) or (6, 2).
That’s fairly straightforward to do:
- Let c_4 and c_6 denote the number of 4's and 6's seen so far.
- For each i from 1 to N,
- If A_i = 2, add c_6 to the answer.
- If A_i = 4, add c_4 to the answer, then increment c_4 by one.
- If A_i = 6, increment c_6 by 1.
- Anything else can be ignored.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Author's code (C++)
#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
template<class T>
using oset = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
// order_of_key(a) -> gives index of the element(number of elements smaller than a)
// find_by_order(a) -> gives the element at index a
#define accelerate ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL)
#define int long long int
#define ld long double
#define mod1 998244353
#define endl "\n"
#define ff first
#define ss second
#define all(x) (x).begin(),(x).end()
#define ra(arr,n) vector<int> arr(n);for(int in = 0; in < n; in++) {cin >> arr[in];}
const int mod = 1e9 + 7;
const int inf = 1e18;
int MOD(int x) {int a1 = (x % mod); if (a1 < 0) {a1 += mod;} return a1;}
int N = 2e5;
vector<int>yep(N + 1, true);
int power( int a, int b) {
int p = 1; while (b > 0) {if (b & 1)p = (p * a); a = (a * a ) ; b >>= 1;}
return p ;
}
void lessgoo()
{
int n;
cin >> n;
ra(arr, n);
int a = 0, ans = 0, f = 0;
for (int i = 0; i < n; i++) {
if (arr[i] == 4)f++;
else if (arr[i] == 6)a++;
else if (arr[i] == 2) {
ans += a;
}
}
cout << ans + f*(f - 1) / 2 << endl;
}
signed main()
{
accelerate;
#ifndef ONLINE_JUDGE
freopen("input.txt", "r", stdin);
freopen("output.txt", "w", stdout);
#endif
int test = 1;
cin >> test ;
for (int tcase = 1; tcase <= test; tcase++)
{
// cout << "Case #" << tcase << ": ";
lessgoo();
}
return 0;
}
Editorialist's code (Python)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
c4, c6 = 0, 0
for x in a:
if x == 4:
ans += c4
c4 += 1
if x == 2: ans += c6
if x == 6: c6 += 1
print(ans)
