代码

#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 256
#define XB 64
#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 u16 recY[XB + 32], recZ[XB + 32];
static u16 cnt[MAXN + 2], stv[MAXN + 2], curs[MAXN + 2];
static u32 ordx[MAXN + 2];
static u64 QA[MAXN + 2];
static u32 QB[MAXN + 2];
static u64 meta[MAXN + 2];
static u16 maskTab[17][16] __attribute__((aligned(32)));
static u32 maskBits[17];
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); }
AVX2 static inline u32 hsum8(__m256i s) {
  __m128i lo = _mm256_castsi256_si128(s), hi = _mm256_extracti128_si256(s, 1);
  __m128i t = _mm_add_epi32(lo, hi);
  t = _mm_add_epi32(t, _mm_shuffle_epi32(t, _MM_SHUFFLE(1, 0, 3, 2)));
  t = _mm_add_epi32(t, _mm_shuffle_epi32(t, _MM_SHUFFLE(2, 3, 0, 1)));
  return (u32)_mm_cvtsi128_si32(t);
}

AVX2 static inline u32 scanv(const u16 *v, int off, int thr) {
  if (off <= 0) return 0;
  const __m256i t = _mm256_set1_epi16((short)thr);
  u32 tot = 0;
  /* CONSTANT TRIP COUNT: the ranK-domain block is exactly BB = 256 = 16 vectors of 16 u16,
     so reading 16 vectors is IN BOUNDS BY CONSTRUCTION (off <= 255), and the exit branch
     that used to depend on off is gone.  Lanes >= off are masked off, so the result is
     identical.  The mask is the artifact's own maskTab idiom. */
  for (int r = 0; r < 16; r++) {
    int k = off - (r << 4);
    k = (k < 0) ? 0 : k; if (k > 16) k = 16;
    __m256i c = _mm256_cmpgt_epi16(t, ldu256(v + (r << 4)));
    c = _mm256_and_si256(c, ldu256(maskTab[k]));
    tot += (u32)__builtin_popcount((u32)_mm256_movemask_epi8(c));
  }
  return tot >> 1;
}

AVX2 static inline u32 scanrec(int nrec, u32 ky, u32 kz) {
  const __m256i ty = _mm256_set1_epi16((short)ky);
  const __m256i tz = _mm256_set1_epi16((short)kz);
  u32 tot = 0;
  const u16 *py = recY, *pz = recZ;
  for (int r = 0; r < (XB >> 4); r++, py += 16, pz += 16) {
    __m256i a = _mm256_cmpgt_epi16(ty, ldu256(py));
    __m256i b = _mm256_cmpgt_epi16(tz, ldu256(pz));
    tot += (u32)__builtin_popcount((u32)_mm256_movemask_epi8(_mm256_and_si256(a, b)));
  }
  return tot >> 1;
}

AVX2 static inline void padrec(void) {
  for (int k = 0; k < XB + 16; k++) { recY[k] = 0x7F7F; recZ[k] = 0x7F7F; }
}

template <int TS> AVX2 static void rebuild_cp(int G) {
  const int S = (TS > 0) ? TS : gS;
  const __m256i bcastLast = _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);
  int fy = 0;
  for (; fy + 1 < G; fy += 2) {
    const u16 *h0 = H + (size_t)fy * S, *h1 = h0 + S;
    u16 *c0 = CP + (size_t)fy * S, *c1 = c0 + S;
    __m256i k0 = _mm256_setzero_si256(), k1 = _mm256_setzero_si256();
    for (int j = 0; j < S; j += 16) {
      __m256i v0 = ldu256(h0 + j), v1 = ldu256(h1 + j);
      v0 = _mm256_add_epi16(v0, _mm256_slli_si256(v0, 2));
      v1 = _mm256_add_epi16(v1, _mm256_slli_si256(v1, 2));
      v0 = _mm256_add_epi16(v0, _mm256_slli_si256(v0, 4));
      v1 = _mm256_add_epi16(v1, _mm256_slli_si256(v1, 4));
      v0 = _mm256_add_epi16(v0, _mm256_slli_si256(v0, 8));
      v1 = _mm256_add_epi16(v1, _mm256_slli_si256(v1, 8));
      v0 = _mm256_add_epi16(v0, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v0, v0, 0x08), bcastLast));
      v1 = _mm256_add_epi16(v1, _mm256_shuffle_epi8(_mm256_permute2x128_si256(v1, v1, 0x08), bcastLast));
      __m256i r0 = _mm256_add_epi16(v0, k0), r1 = _mm256_add_epi16(v1, k1);
      stu256(c0 + j, r0);
      stu256(c1 + j, r1);
      k0 = _mm256_shuffle_epi8(_mm256_permute2x128_si256(r0, r0, 0x11), bcastLast);
      k1 = _mm256_shuffle_epi8(_mm256_permute2x128_si256(r1, r1, 0x11), bcastLast);
    }
  }
  for (; fy < G; fy++) {
    const u16 *h = H + (size_t)fy * S;
    u16 *c = CP + (size_t)fy * S;
    __m256i k = _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), bcastLast));
      __m256i r = _mm256_add_epi16(v, k);
      stu256(c + j, r);
      k = _mm256_shuffle_epi8(_mm256_permute2x128_si256(r, r, 0x11), bcastLast);
    }
  }
  for (int f = 1; f < G; f++) {
    const u16 *p = CP + (size_t)(f - 1) * S;
    u16 *c = CP + (size_t)f * S;
    for (int j = 0; j < S; j += 16)
      stu256(c + j, _mm256_add_epi16(ldu256(c + j), ldu256(p + j)));
  }
}

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; }
  for (int i = 0; i < 17; i++) {
    for (int j = 0; j < 16; j++) maskTab[i][j] = (u16)(j < i ? 0xFFFFu : 0u);
    maskBits[i] = (i >= 16) ? 0xFFFFFFFFu : ((1u << (2 * i)) - 1u);
  }

  /* ---- y and z ranks in one fused pass ---- */
  memset(cnt, 0, (size_t)(n + 1) * 2);
  for (int i = 0; i < n; i++) cnt[y[i]]++;
  { 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 i = 0; i < n; i++) { u32 v = y[i]; meta[i] = ((u64)stv[v] << 48) | (u64)curs[v]++; }

  memset(cnt, 0, (size_t)(n + 1) * 2);
  for (int i = 0; i < n; i++) cnt[z[i]]++;
  { 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 i = 0; i < n; i++) { u32 v = z[i]; meta[i] |= ((u64)stv[v] << 32) | ((u64)curs[v]++ << 16); }

  /* ---- x order ---- */
  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]; stv[v] = (u16)a; curs[v] = (u16)a; a += c; } }
  for (int i = 0; i < n; i++) { u32 v = x[i]; ordx[curs[v]++] = i; }
  for (int e = 0; e < n; e++) {
    const int i = ordx[e];
    const u64 m = meta[i];
    unsigned xv = x[i];
    if (xv > 0x3FFF) xv = 0x3FFF;
    QA[e] = ((m >> 48) << 42) | ((((m >> 32) & 0xFFFF)) << 28) | ((u64)i << 14) | (u64)xv;
    QB[e] = (u32)((m & 0xFFFF) | (((m >> 16) & 0xFFFF) << 14));
  }

  /* ---- 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 u32 xv = (u32)(QA[p] & 0x3FFF);
    while (q < n && ((u32)(QA[q] & 0x3FFF)) == xv) q++;
    for (int e = p; e < q; e++) {
      const u64 m = QA[e];
      const u32 ky = (u32)(m >> 42), kz = (u32)((m >> 28) & 0x3FFF);
      const int FY = (int)BLKIDX(ky), FZ = (int)BLKIDX(kz);
      u32 c = (FY > 0 && FZ > 0) ? CP[(size_t)(FY - 1) * S + (FZ - 1)] : 0;
      c += scanv(YZ + ((size_t)BLK0(FY)), (int)BLOFF(ky), (int)kz);
      c += scanv(ZY + (size_t)BLK0(FZ), (int)BLOFF(kz), BLK0(FY));
      c += scanrec(nrec, (u32)BLK0(FY), (u32)BLK0(FZ));
      out[(m >> 14) & 0x3FFF] = c;
    }
    for (int e = p; e < q; e++) {
      const u32 b = QB[e];
      const u32 ry = b & 0x3FFF, rz = b >> 14;
      H[(size_t)BLKIDX(ry) * S + BLKIDX(rz)]++;
      YZ[ry] = (u16)rz;
      ZY[rz] = (u16)ry;
      recY[nrec] = (u16)ry;
      recZ[nrec] = (u16)rz;
      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<40, 48>(n, x, y, z, out);
  else            engine_t<-1, -1>(n, x, y, z, out);
}

评测结果