Submission #1250110

#TimeUsernameProblemLanguageResultExecution timeMemory
1250110LucaLucaMTriple Peaks (IOI25_triples)C++20
51 / 100
124 ms19372 KiB
#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();

  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 = (jminus + jplus) / 2;
      // a[k] = j - i => i = j - a[k]
      int i = j - a[k];
      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++;
      }
    }
  }

  return answer;
}

std::vector<int> construct_range(int m, int k) {
  return std::vector<int>(m, 0);
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...