Submission #882999

#	Time	Username	Problem	Language	Result	Execution time	Memory
882999		rxlfd314	Pairs (IOI07_pairs)	C++17	100 / 100	1779 ms	479504 KiB

pairs

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ari2 = array<int, 2>;
using ari3 = array<int, 3>;

#define vt vector
#define size(x) (int((x).size()))
#define all(x) begin(x), end(x)

#define REP(a, b, c, d) for (int a = (b); (d) > 0 ? a <= (c) : a >= (c); a += (d))
#define FOR(a, b, c) REP(a, b, c, 1)
#define ROF(a, b, c) REP(a, b, c, -1)

#define chmax(a, b) (a = max(1ll * a, 1ll * (b)))
#define chmin(a, b) (a = min(1ll * a, 1ll * (b)))

struct ST {
  int val;
  ST *lft, *rht;
  ST() : val(0), lft(nullptr), rht(nullptr) {}
  void upd(int p, int v, int tl, int tr) {
    if (tl == tr) {
      val += v;
      return;
    }
    int tm = tl + tr >> 1;
    if (p <= tm) {
      if (!lft)
        lft = new ST();
      lft->upd(p, v, tl, tm);
    } else {
      if (!rht)
        rht = new ST();
      rht->upd(p, v, tm+1, tr);
    }
    val = (lft ? lft->val : 0) + (rht ? rht->val : 0);
  }
  int qry(int ql, int qr, int tl, int tr) {
    if (tl > qr || tr < ql)
      return 0;
    if (ql <= tl && tr <= qr)
      return val;
    int tm = tl + tr >> 1;
    return (lft ? lft->qry(ql, qr, tl, tm) : 0) + (rht ? rht->qry(ql, qr, tm+1, tr) : 0);
  }
};

struct BIT1 {
  vt<ST> fen;
  const int tl, tr;
  BIT1(int n, int l, int r) : fen(vt<ST>(n, ST())), tl(l), tr(r) {}
  void upd(int i, int j, int v) {
    for (; i < size(fen); i += i+1 & -i-1)
      fen[i].upd(j, v, tl, tr);
  }
  int qry(int i, int ql, int qr) {
    int ret = 0;
    for (; ~i; i -= i+1 & -i-1)
      ret += fen[i].qry(ql, qr, tl, tr);
    return ret;
  }
  int qry(int a, int b, int ql, int qr) {
    return qry(b, ql, qr) - qry(a-1, ql, qr);
  }
};

struct BIT2 {
  vt<vt<vt<int>>> fen;
  BIT2(int n, int m, int o) : fen(vt<vt<vt<int>>>(n, vt<vt<int>>(m, vt<int>(o)))) {}
  void upd(int a, int b, int c, int v) {
    for (int i = a; i < size(fen); i += i+1 & -i-1)
      for (int j = b; j < size(fen[i]); j += j+1 & -j-1)
        for (int k = c; k < size(fen[i][j]); k += k+1 & -k-1)
          fen[i][j][k] += v;
  }
  int qry(int a, int b, int c) {
    int ret = 0;
    for (int i = a; ~i; i -= i+1 & -i-1)
      for (int j = b; ~j; j -= j+1 & -j-1)
        for (int k = c; ~k; k -= k+1 & -k-1)
          ret += fen[i][j][k];
    return ret;
  }
  int qry(int a, int d, int b, int e, int c, int f) {
    a--, b--, c--;
    return qry(d, e, f) - qry(a, e, f) - qry(d, b, f) - qry(d, e, c) + qry(a, b, f) + qry(a, e, c) + qry(d, b, c) - qry(a, b, c);
  }
};

signed main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  int B, N, D, M;
  cin >> B >> N >> D >> M;
  if (B == 1) {
    int A[N];
    FOR(i, 0, N-1)
      cin >> A[i];
    sort(A, A+N);
    ll ans = 0;
    FOR(i, 0, N-1)
      ans += upper_bound(A, A+N, A[i]+D) - A - i - 1;
    cout << ans << '\n';
  } else if (B == 2) {
    vt<ari2> A(N);
    for (auto &[a, b] : A)
      cin >> a >> b, a--, b--;
    sort(all(A));
    
    ll ans = 0;
    BIT1 l_fen(M, 0, 2*M-2), r_fen(M, 0, 2*M-2);
    for (auto &[x, y] : A) {
      int l = l_fen.qry(0, y, max(0, x+y-D), 2*M-2);
      int r = y+1 < M ? r_fen.qry(y+1, M-1, max(0, x-y-D+M-1), 2*M-2) : 0;
      ans += l + r;
      //cout << x << ' ' << y << ": " << l << ' ' << r << '\n';
      l_fen.upd(y, x+y, 1);
      r_fen.upd(y, x-y+M-1, 1);
    }
    cout << ans << '\n';
  } else {
    vt<ari3> A(N);
    for (auto &[x, y, z] : A)
      cin >> x >> y >> z, x--, y--, z--;
    sort(all(A));
    
    ll ans = 0;
    // ll: x_i + y_i + z_i - D <= x_j + y_j + z_j
    // lr: x_i + y_i - z_i - D + M-1 <= x_j + y_j - z_j + M-1
    // rl: x_i - y_i + z_i - D + M-1 <= x_j - y_j + z_j + M-1
    // rr: x_i - y_i - z_i - D + 2*M-2 <= x_j - y_j - z_j + 2*M-2
    BIT2 ll_fen(M, M, 3*M-2), lr_fen(M, M, 3*M-2), rl_fen(M, M, 3*M-2), rr_fen(M, M, 3*M-2);
    for (auto &[x, y, z] : A) {
      int bef = ans;
      ans += ll_fen.qry(0, y, 0, z, max(x+y+z-D, 0), 3*M-3);
      ans += z+1 < M ? lr_fen.qry(0, y, z+1, M-1, max(x+y-z-D+M-1, 0), 3*M-3) : 0;
      if (y+1 < M) {
        ans += rl_fen.qry(y+1, M-1, 0, z, max(x-y+z-D+M-1, 0), 3*M-3);
        ans += z+1 < M ? rr_fen.qry(y+1, M-1, z+1, M-1, max(x-y-z-D+2*M-2, 0), 3*M-3) : 0;
      }
      //cout << ans - bef << '\n';
      ll_fen.upd(y, z, x+y+z, 1);
      lr_fen.upd(y, z, x+y-z+M-1, 1);
      rl_fen.upd(y, z, x-y+z+M-1, 1);
      rr_fen.upd(y, z, x-y-z+2*M-2, 1);
    }
    cout << ans << '\n';
  }
} 

Compilation message (stderr)

pairs.cpp: In member function 'void ST::upd(int, int, int, int)':
pairs.cpp:27:17: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   27 |     int tm = tl + tr >> 1;
      |              ~~~^~~~
pairs.cpp: In member function 'int ST::qry(int, int, int, int)':
pairs.cpp:44:17: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   44 |     int tm = tl + tr >> 1;
      |              ~~~^~~~
pairs.cpp: In member function 'void BIT1::upd(int, int, int)':
pairs.cpp:54:33: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   54 |     for (; i < size(fen); i += i+1 & -i-1)
      |                                ~^~
pairs.cpp: In member function 'int BIT1::qry(int, int, int)':
pairs.cpp:59:22: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   59 |     for (; ~i; i -= i+1 & -i-1)
      |                     ~^~
pairs.cpp: In member function 'void BIT2::upd(int, int, int, int)':
pairs.cpp:72:42: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   72 |     for (int i = a; i < size(fen); i += i+1 & -i-1)
      |                                         ~^~
pairs.cpp:73:47: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   73 |       for (int j = b; j < size(fen[i]); j += j+1 & -j-1)
      |                                              ~^~
pairs.cpp:74:52: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   74 |         for (int k = c; k < size(fen[i][j]); k += k+1 & -k-1)
      |                                                   ~^~
pairs.cpp: In member function 'int BIT2::qry(int, int, int)':
pairs.cpp:79:31: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   79 |     for (int i = a; ~i; i -= i+1 & -i-1)
      |                              ~^~
pairs.cpp:80:33: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   80 |       for (int j = b; ~j; j -= j+1 & -j-1)
      |                                ~^~
pairs.cpp:81:35: warning: suggest parentheses around '+' in operand of '&' [-Wparentheses]
   81 |         for (int k = c; ~k; k -= k+1 & -k-1)
      |                                  ~^~
pairs.cpp: In function 'int main()':
pairs.cpp:135:11: warning: unused variable 'bef' [-Wunused-variable]
  135 |       int bef = ans;
      |           ^~~

• Subtask 1: 6 / 6
• Subtask 2: 6 / 6
• Subtask 3: 6 / 6
• Subtask 4: 6 / 6
• Subtask 5: 6 / 6
• Subtask 6: 7 / 7
• Subtask 7: 7 / 7
• Subtask 8: 8 / 8
• Subtask 9: 8 / 8
• Subtask 10: 10 / 10
• Subtask 11: 10 / 10
• Subtask 12: 10 / 10
• Subtask 13: 10 / 10