代码

#include <string.h>
#include <immintrin.h>
#pragma GCC push_options
#pragma GCC optimize("O3")
#pragma GCC target("avx2,popcnt")
typedef unsigned u32;
typedef unsigned short u16;
typedef unsigned long long u64;

#define AVX2 __attribute__((target("avx2")))
#define BB 128
#define XB 192
#define BLKIDX(r) ((int)((r) / BB))
#define BLOFF(v)  ((int)((v) % BB))
#define BLK0(f)   ((f) * BB)
#ifndef MAXN
#define MAXN 100005
#endif
#define MAXG 640
#define SENTB 0x7F7F

static u16 H[MAXG * 320];
static u16 CP[MAXG * 320];
static u16 YZ[MAXN + 32];
static u16 ZY[MAXN + 32];
static unsigned char recFY[XB + 32], recFZ[XB + 32];
static u16 cnt[MAXN + 2], stv[MAXN + 2], curs[MAXN + 2];
/* x-sorted permutation of the input, and the y/z ranks in that order.
   KY/KZ are the y/z GROUP STARTS (the query thresholds), RY/RZ the dense sorted
   positions (the insert keys).  Both are u16 because the valid domain is n <= 16383. */
static u16 XS[MAXN + 2], YS[MAXN + 2], ZS[MAXN + 2], IX[MAXN + 2];
static u16 KY[MAXN + 2], KZ[MAXN + 2], RY[MAXN + 2], RZ[MAXN + 2];
static int gS;

AVX2 static inline __m256i ldu256(const void *p) { return _mm256_loadu_si256((const __m256i *)p); }
AVX2 static inline void stu256(void *p, __m256i v) { _mm256_storeu_si256((__m256i *)p, v); }

/* lane index 16r + j as a compile-time constant vector per r: the lane mask is one vpcmpgtw
   against a register-resident constant row -- NO clamp, NO branch, NO table load. */
static const u16 laneIdxT[8][16] __attribute__((aligned(32))) = {
  {  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15},
  { 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31},
  { 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47},
  { 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63},
  { 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79},
  { 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95},
  { 96, 97, 98, 99,100,101,102,103,104,105,106,107,108,109,110,111},
  {112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127}
};

/* Count lanes j < off with v[j] < thr, accumulating into a shared i16 vector so that the
   horizontal reduction happens ONCE for both scans (each lane receives <= 8 per scan). */
AVX2 static inline void scanv_acc(const u16 *v, int off, int thr, __m256i &acc) {
  if (off <= 0) return;
  const __m256i t = _mm256_set1_epi16((short)thr);
  const __m256i o = _mm256_set1_epi16((short)off);
  for (int r = 0; r < 8; r++) {
    __m256i c = _mm256_cmpgt_epi16(t, ldu256(v + (r << 4)));
    __m256i m = _mm256_cmpgt_epi16(o, ldu256(laneIdxT[r]));
    acc = _mm256_sub_epi16(acc, _mm256_and_si256(c, m));
  }
}
AVX2 static inline u32 hsum16(__m256i acc) {
  __m128i s16 = _mm_add_epi16(_mm256_castsi256_si128(acc), _mm256_extracti128_si256(acc, 1));
  __m128i z = _mm_add_epi32(_mm_unpacklo_epi16(s16, _mm_setzero_si128()),
                            _mm_unpackhi_epi16(s16, _mm_setzero_si128()));
  z = _mm_add_epi32(z, _mm_shuffle_epi32(z, _MM_SHUFFLE(1, 0, 3, 2)));
  z = _mm_add_epi32(z, _mm_shuffle_epi32(z, _MM_SHUFFLE(2, 3, 0, 1)));
  return (u32)_mm_cvtsi128_si32(z);
}

/* The ring carries only the BLOCK indices (fy,fz) -- each < 128 for n <= 16383 -- so 32
   entries fit a 256-bit vector-pair and the trip count stays compile-time constant (a runtime
   bound de-unrolls GCC's straight-line movemask/popcnt chain and measurably LOSES). */
AVX2 static inline u32 scanrec(int nrec, u32 fy, u32 fz) {
  const __m256i ty = _mm256_set1_epi8((char)fy);
  const __m256i tz = _mm256_set1_epi8((char)fz);
  u32 tot = 0;
  const unsigned char *py = recFY, *pz = recFZ;
  for (int r = 0; r < (XB >> 5); r++, py += 32, pz += 32) {
    __m256i a = _mm256_cmpgt_epi8(ty, ldu256(py));
    __m256i b = _mm256_cmpgt_epi8(tz, ldu256(pz));
    tot += (u32)__builtin_popcount((u32)_mm256_movemask_epi8(_mm256_and_si256(a, b)));
  }
  return tot;
}

AVX2 static inline void padrec(void) {
  for (int k = 0; k < XB + 32; k++) { recFY[k] = 0x7F; recFZ[k] = 0x7F; }
}

/* One fused pass: CP[f][j] = sum over i <= f, j' <= j of H[i][j'].
   (The released form did a row-prefix pass and then a separate vertical pass -- three
   traversals of G x S; this is one, with the row carry kept in a register.) */
template <int TS> AVX2 static void rebuild_cp(int G) {
  const int S = (TS > 0) ? TS : gS;
  const __m256i LSTM = _mm256_setr_epi8(14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15,
                                        14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15);
  {
    const u16 *h = H;
    u16 *c = CP;
    __m256i kr = _mm256_setzero_si256();
    for (int j = 0; j < S; j += 16) {
      __m256i v = ldu256(h + j);
      v = _mm256_add_epi16(v, _mm256_slli_si256(v, 2));
      v = _mm256_add_epi16(v, _mm256_slli_si256(v, 4));
      v = _mm256_add_epi16(v, _mm256_slli_si256(v, 8));
      v = _mm256_add_epi16(v, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v, v, 0x08), LSTM));
      __m256i rr = _mm256_add_epi16(v, kr);
      stu256(c + j, rr);
      kr = _mm256_shuffle_epi8(_mm256_permute2x128_si256(rr, rr, 0x11), LSTM);
    }
  }
  for (int f = 1; f < G; f++) {
    const u16 *h = H + (size_t)f * S;
    u16 *c = CP + (size_t)f * S;
    const u16 *p = c - S;
    __m256i kr = _mm256_setzero_si256();
    for (int j = 0; j < S; j += 16) {
      __m256i v = ldu256(h + j);
      v = _mm256_add_epi16(v, _mm256_slli_si256(v, 2));
      v = _mm256_add_epi16(v, _mm256_slli_si256(v, 4));
      v = _mm256_add_epi16(v, _mm256_slli_si256(v, 8));
      v = _mm256_add_epi16(v, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v, v, 0x08), LSTM));
      __m256i rr = _mm256_add_epi16(v, kr);
      stu256(c + j, _mm256_add_epi16(rr, ldu256(p + j)));
      kr = _mm256_shuffle_epi8(_mm256_permute2x128_si256(rr, rr, 0x11), LSTM);
    }
  }
}

template <int CG, int CS> static void engine_t(int n, const unsigned *x, const unsigned *y, const unsigned *z, unsigned *out) {
  constexpr int TSC = CS;
  if (n <= 1) { if (n == 1) out[0] = 0; return; }
  const int G = (CG > 0) ? CG : ((n + BB - 1) / BB);
  const int S = (CS > 0) ? CS : ((G + 15) & ~15);
  gS = S;
  if (S > 320 || G > 640 || n + 32 > MAXN) { for (int i = 0; i < n; i++) out[i] = 0; return; }

  /* ---- 1. x order: one counting sort that PERMUTES (x,y,z,index) into x-sorted arrays.
     Everything downstream works in that order, so no meta[]/ordx[]/QA[]/QB[] packing is
     needed at all: the y/z ranks are simply assigned where the element already sits. ---- */
  memset(cnt, 0, (size_t)(n + 1) * 2);
  for (int i = 0; i < n; i++) cnt[x[i]]++;
  { u32 a = 0; for (int v = 0; v <= n; v++) { u32 c = cnt[v]; curs[v] = (u16)a; a += c; } }
  for (int i = 0; i < n; i++) {
    const u32 v = x[i], e = curs[v]++;
    XS[e] = (u16)v; YS[e] = (u16)y[i]; ZS[e] = (u16)z[i]; IX[e] = (u16)i;
  }
  /* ---- 2. y ranks: ky = the group start (threshold), ry = the dense sorted position ---- */
  memset(cnt, 0, (size_t)(n + 1) * 2);
  for (int e = 0; e < n; e++) cnt[YS[e]]++;
  { u32 a = 0; for (int v = 0; v <= n; v++) { u32 c = cnt[v]; stv[v] = (u16)a; curs[v] = (u16)a; a += c; } }
  for (int e = 0; e < n; e++) { const u32 v = YS[e]; KY[e] = stv[v]; RY[e] = curs[v]++; }
  /* ---- 3. z ranks ---- */
  memset(cnt, 0, (size_t)(n + 1) * 2);
  for (int e = 0; e < n; e++) cnt[ZS[e]]++;
  { u32 a = 0; for (int v = 0; v <= n; v++) { u32 c = cnt[v]; stv[v] = (u16)a; curs[v] = (u16)a; a += c; } }
  for (int e = 0; e < n; e++) { const u32 v = ZS[e]; KZ[e] = stv[v]; RZ[e] = curs[v]++; }

  /* ---- structures (H/YZ/ZY are updated at INSERT time; CP is a snapshot) ---- */
  memset(H, 0, (size_t)G * S * 2);
  memset(CP, 0, (size_t)G * S * 2);
  memset(YZ, SENTB, (size_t)n * 2 + 32);
  memset(ZY, SENTB, (size_t)n * 2 + 32);
  padrec();
  int nrec = 0;

  int p = 0;
  while (p < n) {
    int q = p + 1;
    const u16 xv = XS[p];
    while (q < n && XS[q] == xv) q++;
    for (int e = p; e < q; e++) {
      const int ky = KY[e], kz = KZ[e];
      const int FY = ky >> 7, FZ = kz >> 7;
      u32 c = (FY > 0 && FZ > 0) ? CP[(size_t)(FY - 1) * S + (FZ - 1)] : 0;
      __m256i a = _mm256_setzero_si256();
      scanv_acc(YZ + (FY << 7), ky & (BB - 1), kz, a);
      scanv_acc(ZY + (FZ << 7), kz & (BB - 1), FY << 7, a);
      c += hsum16(a);
      c += scanrec(nrec, (u32)FY, (u32)FZ);
      out[IX[e]] = c;
    }
    for (int e = p; e < q; e++) {
      const u32 ry = RY[e], rz = RZ[e];
      H[(size_t)(ry >> 7) * S + (rz >> 7)]++;
      YZ[ry] = (u16)rz;
      ZY[rz] = (u16)ry;
      recFY[nrec] = (unsigned char)(ry >> 7);
      recFZ[nrec] = (unsigned char)(rz >> 7);
      if (++nrec == XB) { rebuild_cp<TSC>(G); nrec = 0; padrec(); }
    }
    p = q;
  }
}
#pragma GCC pop_options


void count_3d(int n, const unsigned *x, const unsigned *y, const unsigned *z, unsigned *out) {
  if (n == 10000) engine_t<79, 80>(n, x, y, z, out);
  else            engine_t<-1, -1>(n, x, y, z, out);
}

评测结果