#include <iostream>
#include <vector>
#include <algorithm>
#include <cassert>
#include <tuple>
#include <utility>
#include <random>
#include <ctime>
#include <iomanip>
#include "triples.h"
std::mt19937 rng(123);
using ll = long long;
#define debug(x) #x << " = " << x << '\n'
const int SQRT = 512;
bool same(std::vector<int> a, std::vector<int> b) {
if ((int) a.size() != (int) b.size()) {
return false;
}
std::sort(a.begin(), a.end());
std::sort(b.begin(), b.end());
for (int i = 0; i < (int) a.size(); i++) {
if (a[i] != b[i]) {
return false;
}
}
return true;
}
ll count_triples(std::vector<int> a) {
int n = (int) a.size();
auto isGood = [&](int i, int j, int k) {
if (!(0 <= i && i < j && j < k && k < n)) {
return false;
}
return same(std::vector<int>{a[i], a[j], a[k]}, std::vector<int>{j - i, k - i, k - j});
};
std::vector<std::tuple<int, int, int>> cand;
for (int i = 0; i < n; i++) {
{ // I. a[i] = k - i
// => k = a[i] + i
int k = a[i] + i;
if (i < k && k < n) {
// a[k] e j - i sau k - j
// a[k] = j - i => j e a[k] + i
// a[k] = k - j => j e k - a[k]
for (int j : {a[k] + i, k - a[k]}) {
if (i < j && j < k) {
if (isGood(i, j, k)) {
cand.push_back({i, j, k});
}
}
}
}
}
{ // II. a[i] = j - i
int j = a[i] + i;
if (i < j && j < n) {
// a[j] e k - i sau k - j
for (int k : {a[j] + i, a[j] + j}) {
if (j < k && k < n) {
if (isGood(i, j, k)) {
cand.push_back({i, j, k});
}
}
}
}
}
}
// mai tratez cazul asta: (a[i], a[j], a[k]) = (k - j, j - i, k - i), aici egalitatea se intelege ca fiind intre cele doua siruri (nu ca multiset uri)
for (int j = 0; j < n; j++) {
// a[j] = j - i
int i = j - a[j];
if (0 <= i && i < j) {
// a[i] = k - j
int k = a[i] + j;
if (isGood(i, j, k)) {
cand.push_back({i, j, k});
}
}
}
// ok acum ramane doar urmatorul caz:
// a[i] = k - j
// a[j] = k - i
// a[k] = j - i
// a[i] = k - j, a[j] = k - i => a[i] - a[j] = -j + i <=> a[i] - i = a[j] - j
std::vector<std::vector<int>> is(2 * n); // js[x] contine toate i urile care au a[i] - i + (n - 1) = x
for (int i = 0; i < n; i++) {
is[a[i] - i + n - 1].push_back(i);
}
std::vector<int> bigf;
std::sort(cand.begin(), cand.end());
cand.erase(std::unique(cand.begin(), cand.end()), cand.end());
int answer = (int) cand.size();
// std::cout << debug((int) cand.size());
for (int i = 0; i < n; i++) {
if ((int) is[a[i] - i + n - 1].size() <= SQRT) {
for (int j : is[a[i] - i + n - 1]) {
// a[i] = k - j => k = a[i] + j
int k = a[i] + j;
if (0 <= i && i < j && j < k && k < n && a[i] == k - j && a[j] == k - i && a[k] == j - i && a[i] != a[k] && a[i] != a[j] && a[j] != a[k]) {
answer++;
}
}
} else {
bigf.push_back(a[i] - i);
}
}
std::sort(bigf.begin(), bigf.end());
bigf.erase(std::unique(bigf.begin(), bigf.end()), bigf.end());
for (int k = 0; k < n; k++) {
for (int delta : bigf) {
// stiu k si delta = a[i] - i = a[j] - j
// a[j] + j = a[k] + k
int jminus = delta;
int jplus = a[k] + k;
int j = (jplus - jminus) / 2;
// if (k == n - 1) {
// std::cout << debug(delta);
// std::cout << debug(j);
// }
// a[k] = j - i => i = j - a[k]
int i = j - a[k];
if ((jminus + jplus) % 2 == 0 && 0 <= i && i < j && j < k && k < n && a[i] == k - j && a[j] == k - i && a[k] == j - i && a[i] != a[k] && a[i] != a[j] && a[j] != a[k]) {
answer++;
}
}
}
return answer;
}
ll brute(std::vector<int> a) {
int n = (int) a.size();
int ret = 0;
auto isGood = [&](int i, int j, int k) {
if (!(0 <= i && i < j && j < k && k < n)) {
return false;
}
return same(std::vector<int>{a[i], a[j], a[k]}, std::vector<int>{j - i, k - i, k - j});
};
for (int i = 0; i + 2 < n; i++) {
for (int j = i + 1; j + 1 < n; j++) {
for (int k = j + 1; k < n; k++) {
if (isGood(i, j, k)) {
ret++;
}
}
}
}
return (ll) ret;
}
std::vector<int> construct_range(int m, int k) {
return std::vector<int>(m, 0);
}
# | 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... |