#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_set>
#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
using namespace std;
int64_t pack(int i, int j, int k) {
return ((int64_t)i << 40) | ((int64_t)j << 20) | ((int64_t)k);
}
void sort3(int &a0, int &a1, int &a2) {
if (a0 > a1) swap(a0, a1);
if (a1 > a2) swap(a1, a2);
if (a0 > a1) swap(a0, a1);
}
int add(int a[], int n, int v) {
#pragma GCC unroll 8
for (int i = 0; i < n; ++i)
if (a[i] == v) return n;
a[n++] = v; return n;
}
void check(int i, int j, int k, vector<int>&a \
, unordered_set<int64_t> &s) {
vector<int> x = { j - i, k - i, k - j };
vector<int> y = { a[i], a[j], a[k] };
sort3(x[0], x[1], x[2]);
sort3(y[0], y[1], y[2]);
if((x[0] == y[0]) && (x[1] == y[1]) && (x[2] == y[2]))
s.insert(pack(i, j, k));
}
void findi(int j, int k, vector<int>&a \
, unordered_set<int64_t> &ms) {
if ((k < 0) || ((k >= a.size()))) return;
if ((j < 0) || ((j >= a.size()))) return;
if ((k + a[k]) != (j + a[j])) return;
int op[4], n = 0;
n = add(op, n, j - a[j]);
n = add(op, n, j - a[k]);
n = add(op, n, k - a[j]);
n = add(op, n, k - a[k]);
#pragma GCC unroll 4
for (int it = 0; it < n; ++it) {
int i = op[it];
if ((i >= 0) && (i < a.size()))
check(i, j, k, a, ms);
}
}
void findk(int i, int j, vector<int>& a \
, unordered_set<int64_t> &ms) {
if ((j < 0) || ((j >= a.size()))) return;
if ((i < 0) || ((i >= a.size()))) return;
int op[4], n = 0, sum = j + a[j];
n = add(op, n, j + a[j]);
n = add(op, n, j + a[i]);
n = add(op, n, i + a[j]);
n = add(op, n, i + a[i]);
#pragma GCC unroll 4
for (int it = 0; it < n; ++it) {
int k = op[it];
if ((k >= 0) && (k < a.size()))
if ((k + a[k]) == sum)
check(i, j, k, a, ms);
}
}
long long count_triples(vector<int> a) {
unordered_set<int64_t> ans;
int n = a.size();
ans.reserve(max(1, n * 16));
ans.max_load_factor(0.25f);
for (int x = 0; x < n; ++x)
for (int y : { x - a[x], x + a[x]})
if (!((y < 0) || ((y >= a.size())))) {
vector<int> op = { x - a[x], x + a[x], y + a[y], y - a[y],
y + a[x], y - a[x], x + a[y], x - a[y] };
sort(op.begin(), op.end());
op.erase(unique(op.begin(), op.end()), op.end());
#pragma GCC unroll 8
for (int z : op) {
if ((z < 0) || ((z >= a.size()))) continue;
vector<int> op = { x, y, z };
sort3(op[0], op[1], op[2]);
check(op[0], op[1], op[2], a, ans);
}
}
vector<vector<int>> x(n + n), y(n + n);
for (int i = 0; i < n; ++i)
x[a[i] - i + n].push_back(i);
for (int i = 0; i < n; ++i)
y[a[i] + i].push_back(i);
for (int j = 0; j < n; ++j) {
auto op1 = x[a[j] - j + n], op2 = y[a[j] + j];
if (op1.size() < op2.size())
for (int i : op1)
findk(i, j, a, ans);
else
for (int k : op1)
findi(j, k, a, ans);
}
return (int)ans.size();
}
vector<int> construct_range(int M, int K) {
vector<int> a = { 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 }; return a;
}
//int main() { vector<int> a = { 4, 1, 4, 3, 2, 6, 1 }; cout << count_triples(a); }
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |