Submission #1294215

#TimeUsernameProblemLanguageResultExecution timeMemory
1294215avighnaPairs (IOI07_pairs)C++20
80 / 100
4091 ms197404 KiB
#pragma GCC optimize("Ofast", "unroll-all-loops")

#include <algorithm>
#include <iostream>
#include <unordered_map>
#include <vector>

template <typename T, int d>
class fenwick_tree;
template <typename T>
class fenwick_tree<T, 1> {
private:
  std::size_t n;
  std::unordered_map<std::size_t, T> f;

public:
  fenwick_tree(std::size_t n) : n(n) {}
  void apply(std::size_t i, T x) {
    for (++i; i <= n; i += i & -i) {
      f[i] = f[i] + x;
    }
  }
  T query(std::size_t i) {
    T ans = T{};
    for (++i; i > 0; i -= i & -i) {
      if (f.contains(i)) {
        ans = ans + f.at(i);
      }
    }
    return ans;
  }
  T query(std::size_t l, std::size_t r)
    requires requires(T a, T b) { a - b; }
  {
    return l == 0 ? query(r) : query(r) - query(l - 1);
  }
};
template <typename T, int d>
class fenwick_tree {
private:
  std::size_t n;
  std::vector<fenwick_tree<T, d - 1>> f;

public:
  template <typename... ints>
  fenwick_tree(std::size_t n, ints... dims) : n(n), f(n + 1, fenwick_tree<T, d - 1>{dims...}) {}
  template <typename... ints>
  void apply(std::size_t i, ints... next) {
    for (++i; i <= n; i += i & -i) {
      f[i].apply(next...);
    }
  }
  template <typename... ints>
    requires(sizeof...(ints) == d - 1)
  T query(std::size_t i, ints... dims) {
    T ans = T{};
    for (++i; i > 0; i -= i & -i) {
      ans = ans + f[i].query(dims...);
    }
    return ans;
  }
  template <typename... ints>
  T query(std::size_t l, std::size_t r, ints... dims)
    requires requires(T a, T b) { a - b; }
  {
    T ans = T{};
    for (++r; r > 0; r -= r & -r) {
      ans = ans + f[r].query(dims...);
    }
    for (; l > 0; l -= l & -l) {
      ans = ans - f[l].query(dims...);
    }
    return ans;
  }
};

using namespace std;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);

  int b, n, d, m;
  cin >> b >> n >> d >> m;
  m++;
  int64_t ans = 0;
  if (b == 1) {
    vector<int> a(n);
    for (int &i : a) {
      cin >> i;
    }
    sort(a.begin(), a.end());
    for (int i = 0, j = 0; i < n; ++i) {
      while (j < n && a[j] - a[i] <= d) {
        j++;
      }
      ans += j - i - 1;
    }
  } else if (b == 2) {
    fenwick_tree<int, 2> fenwick(3 * m, 2 * m);
    for (int i = 0, u, v; i < n; ++i) {
      cin >> u >> v;
      int x = v + u + m, y = v - u + m;
      ans += fenwick.query(max(0, x - d), min(x + d, 3 * m - 1), max(0, y - d), min(y + d, 2 * m - 1));
      fenwick.apply(x, y, 1);
    }
  } else {
    fenwick_tree<int, 3> fenwick(3 * m, 2 * m, m);
    for (int i = 0, u, v, w; i < n; ++i) {
      cin >> u >> v >> w;
      int x = v + u + m, y = v - u + m, z = w;
      for (int dz = -d; dz <= d; ++dz) {
        int dd = d - abs(dz);
        if (dd >= 0 && z + dz >= 0 && z + dz < m) {
          ans += fenwick.query(max(0, x - dd), min(x + dd, 3 * m - 1), max(0, y - dd), min(y + dd, 2 * m - 1), z + dz, z + dz);
        }
      }
      fenwick.apply(x, y, z, 1);
    }
  }

  cout << ans << '\n';
}

Compilation message (stderr)

pairs.cpp: In instantiation of 'fenwick_tree<T, d>::fenwick_tree(std::size_t, ints ...) [with ints = {int}; T = int; int d = 2; std::size_t = long unsigned int]':
pairs.cpp:100:46:   required from here
pairs.cpp:46:85: warning: narrowing conversion of 'dims#0' from 'int' to 'std::size_t' {aka 'long unsigned int'} [-Wnarrowing]
   46 |   fenwick_tree(std::size_t n, ints... dims) : n(n), f(n + 1, fenwick_tree<T, d - 1>{dims...}) {}
      |                                                                                     ^~~~
pairs.cpp: In instantiation of 'fenwick_tree<T, d>::fenwick_tree(std::size_t, ints ...) [with ints = {int, int}; T = int; int d = 3; std::size_t = long unsigned int]':
pairs.cpp:108:49:   required from here
pairs.cpp:46:85: warning: narrowing conversion of 'dims#0' from 'int' to 'std::size_t' {aka 'long unsigned int'} [-Wnarrowing]
#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...