PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Simple
PREREQUISITES:
None
PROBLEM:
Given an array A, count the number of triplets (i, j, k) such that |A_i - A_k| = |i-j|+|j-k|.
It is known that 1 \leq A_i \leq 100.
EXPLANATION:
Note the unusual constraint on A_i: the elements are \leq 100 which is quite small.
We care about the equality
The left side of this is the difference of two array elements, and since they’re both between 1 and 100, their difference cannot exceed 100 either.
So, if |i-j| + |j-k| is larger than 100, such a pair can never be valid.
|i-j| + |j-k| is minimized when j is between i and k, and this minimum value is |i-k|.
Intuitively, |i-j| + |j-k| represents the distance of walking from i to j and then j to k on the number line; and this distance is minimized whenever j lies between i and k.
So, we must have |A_i - A_k| \geq |i - k|; otherwise no valid j can exist.
Once again utilizing the fact that the left side is no more than 100, we see that |i-k| \leq 100 as well.
Let’s fix the index i.
Then, since |i-k| \leq 100, k must lie in the range [i-100, i+100].
There are at most 200 options for k, so we can try them all.
Once i and k are fixed, let’s try to find values of j.
We have a couple of cases.
- If |A_i - A_k| \lt |i - k|, then no valid j can exist, as noted above.
- If |A_i - A_k| = |i - k|, then every j between i and k is valid.
There are |i-k| + 1 choices of j here. - If |A_i - A_k| \gt |i - k|, we need a bit more analysis.
Recall that |i-j| + |j-k| is minimized when j lies between i and k, and equals |i-k|.
Let’s analyze what happens for other values of j.
Then,
- If j \geq \max(i,k), we have |i-j| + |j-k| = (j-i) + (j-k).
We want |A_i - A_k| = 2j - i - k, which can be solved uniquely for j as
- Similarly, considering j \leq \min(i,k) gives another unique value of j, that being
We thus obtain at most two values of j that can possibly be valid.
Check validity for both of them (they must be integers between 1 and N), and add to the answer appropriately.
We check at most 200\cdot N pairs of (i,k), each of which is processed in constant time - which is easily fast enough to run within a second.
TIME COMPLEXITY:
\mathcal{O}(N\cdot \max(A)) per testcase, which is fast here since \max(A) \leq 100.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18
mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());
void Solve()
{
int n; cin >> n;
vector <int> a(n + 1);
for (int i = 1; i <= n; i++){
cin >> a[i];
}
int ans = 0;
for (int i = 1; i <= n; i++){
for (int j = max(i - 100, 1LL); j <= min(i + 100, n); j++){
int v = abs(a[i] - a[j]);
int m1 = min(i, j);
int m2 = max(i, j);
int d = (m2 - m1);
if (d % 2 != v % 2 || v < d) continue;
if (v == d){
ans += (m2 - m1 + 1);
continue;
}
if (v <= abs(1 - i) + abs(1 - j)) ans++;
if (v <= abs(n - i) + abs(n - j)) ans++;
}
}
cout << ans << "\n";
}
int32_t main()
{
auto begin = std::chrono::high_resolution_clock::now();
ios_base::sync_with_stdio(0);
cin.tie(0);
int t = 1;
// freopen("in", "r", stdin);
// freopen("out", "w", stdout);
cin >> t;
for(int i = 1; i <= t; i++)
{
//cout << "Case #" << i << ": ";
Solve();
}
auto end = std::chrono::high_resolution_clock::now();
auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n";
return 0;
}
Tester'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<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)
template<typename T>
void amin(T &a, T b) {
a = min(a,b);
}
template<typename T>
void amax(T &a, T b) {
a = max(a,b);
}
#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif
/*
*/
const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;
void solve(int test_case)
{
ll n; cin >> n;
vector<ll> a(n+5);
rep1(i,n) cin >> a[i];
ll ans = 0;
rep1(i,n){
for(int k = i+1; k <= min(i+100,(int)n); ++k){
ll x = abs(a[i]-a[k]);
// j < i
{
ll v = i+k-x;
if(v >= 0 and (!(v&1))){
v >>= 1;
if(v >= 1 and v <= i-1){
ans++;
}
}
}
// i <= j <= k
{
if(k-i == x){
ans += k-i+1;
}
}
// j > k
{
ll v = i+k+x;
if(v >= 0 and !(v&1)){
v >>= 1;
if(v > k and v <= n){
ans++;
}
}
}
}
}
ans *= 2;
ans += n;
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
return 0;
}
Editorialist's code (PyPy3)
import sys
input = sys.stdin.readline
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
for i in range(n):
for k in range(max(0, i-105), min(n, i+105)):
indd = abs(k - i)
eled = abs(a[i] - a[k])
if indd == eled: ans += abs(k - i) + 1
elif indd < eled:
for j in (i+k+eled, i+k-eled):
if j%2 == 0 and j >= 0 and j < 2*n: ans += 1
print(ans)