代码

// Duck.ac 1001: V23, C++17 / i3-8100 AVX2. Readability-only revision.
// Pipeline: uint32 -> packed 24-bit buckets -> uint16 buckets -> byte columns.
// Single-threaded; static workspace makes this entry point non-reentrant.
// 12-input / 39-comparator network by Bert Dobbelaere:
// https://bertdobbelaere.github.io/sorting_networks.html#N12L39D9
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <cstdlib>
#include <immintrin.h>
#include <utility>
#pragma GCC optimize("O3,unroll-loops,no-strict-aliasing")
#pragma GCC target("avx2,bmi,bmi2,popcnt,lzcnt")
static_assert(sizeof(unsigned) == 4 && sizeof(uint16_t) == 2, "32-bit unsigned required");

namespace fastsort {
constexpr unsigned kRadix = 256, kRecordsPerBlock = 84, kBlockBytes = 256;
constexpr unsigned kLeafLimit = 8192;
constexpr int kSmallSort = 4096;
constexpr size_t kScratchOffset = 380'000'000, kWorkspaceBytes = 384'000'000;
using U32 = unsigned; using Byte = uint8_t; using U16 = uint16_t; using U64 = uint64_t;
using Vec256 = __m256i; using Vec128 = __m128i;
template <class T> inline T load(const void *src) { T value; std::memcpy(&value, src, sizeof value); return value; }
template <class F, size_t... I> [[gnu::always_inline]] inline void each(F &&f, std::index_sequence<I...>) {
  (f(std::integral_constant<size_t, I>{}), ...);
}
template <size_t N, class F> [[gnu::always_inline]] inline void repeat(F &&f) {
  each(std::forward<F>(f), std::make_index_sequence<N>{});
}
} // namespace fastsort

namespace byte_sort {
using namespace fastsort;
[[gnu::noinline]] inline void counting_sort(const U16 *src, U32 n, U32 *dst, U32 base) {
  alignas(64) U32 counts[65536] = {};
  for (U32 i = 0; i < n; ++i) ++counts[load<U16>(src + i)];
  for (U32 v = 0; v < 65536; ++v)
    for (U32 c = counts[v]; c; --c) *dst++ = base | v;
}
inline void expand8(Vec128 x, U32 *dst, U32 c, U32 base, U32 *end) {
  Vec256 y = _mm256_or_si256(_mm256_cvtepu8_epi32(x), _mm256_set1_epi32(base));
  if (end - dst >= 8) _mm256_storeu_si256((Vec256 *)dst, y);
  else _mm256_maskstore_epi32((int *)dst,
      _mm256_cmpgt_epi32(_mm256_set1_epi32(c), _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7)), y);
}
// Sort 32 independent columns, then transpose each into an 8-byte head and a 4-byte tail.
[[gnu::always_inline]] inline void sort12_columns(const Byte *src, Byte *out, Byte *tail) {
  Vec256 x[12], a[8], b[8];
  repeat<12>([&](auto j) __attribute__((always_inline)) {
    x[j] = _mm256_load_si256((const Vec256 *)(src + 256 * j));
  });
  // clang-format off
#define CMP(A, B) do { Vec256 lo = _mm256_min_epu8(x[A], x[B]); \
  x[B] = _mm256_max_epu8(x[A], x[B]); x[A] = lo; } while (0)
  CMP(0,8); CMP(1,7); CMP(2,6); CMP(3,11); CMP(4,10); CMP(5,9);
  CMP(0,1); CMP(2,5); CMP(3,4); CMP(6,9); CMP(7,8); CMP(10,11);
  CMP(0,2); CMP(1,6); CMP(5,10); CMP(9,11);
  CMP(0,3); CMP(1,2); CMP(4,6); CMP(5,7); CMP(8,11); CMP(9,10);
  CMP(1,4); CMP(3,5); CMP(6,8); CMP(7,10);
  CMP(1,3); CMP(2,5); CMP(6,9); CMP(8,10);
  CMP(2,3); CMP(4,5); CMP(6,7); CMP(8,9);
  CMP(4,6); CMP(5,7);
  CMP(3,4); CMP(5,6); CMP(7,8);
  // clang-format on
#undef CMP
#pragma GCC unroll 4
  for (U32 i = 0; i < 4; ++i) {
    a[2 * i] = _mm256_unpacklo_epi8(x[2 * i], x[2 * i + 1]);
    a[2 * i + 1] = _mm256_unpackhi_epi8(x[2 * i], x[2 * i + 1]);
  }
#pragma GCC unroll 2
  for (U32 i = 0; i < 2; ++i) {
#pragma GCC unroll 2
    for (U32 j = 0; j < 2; ++j) {
      b[4 * i + 2 * j] = _mm256_unpacklo_epi16(a[4 * i + j], a[4 * i + j + 2]);
      b[4 * i + 2 * j + 1] = _mm256_unpackhi_epi16(a[4 * i + j], a[4 * i + j + 2]);
    }
  }
#pragma GCC unroll 4
  for (U32 i = 0; i < 4; ++i) {
    Vec256 lo = _mm256_unpacklo_epi32(b[i], b[i + 4]);
    Vec256 hi = _mm256_unpackhi_epi32(b[i], b[i + 4]);
    _mm256_store_si256((Vec256 *)(out + 32 * i), _mm256_permute2x128_si256(lo, hi, 0x20));
    _mm256_store_si256((Vec256 *)(out + 128 + 32 * i), _mm256_permute2x128_si256(lo, hi, 0x31));
  }
  a[0] = _mm256_unpacklo_epi8(x[8], x[9]);
  a[1] = _mm256_unpackhi_epi8(x[8], x[9]);
  a[2] = _mm256_unpacklo_epi8(x[10], x[11]);
  a[3] = _mm256_unpackhi_epi8(x[10], x[11]);
  b[0] = _mm256_unpacklo_epi16(a[0], a[2]);
  b[1] = _mm256_unpackhi_epi16(a[0], a[2]);
  b[2] = _mm256_unpacklo_epi16(a[1], a[3]);
  b[3] = _mm256_unpackhi_epi16(a[1], a[3]);
#pragma GCC unroll 2
  for (U32 i = 0; i < 2; ++i) {
    _mm256_store_si256((Vec256 *)(tail + 32 * i), _mm256_permute2x128_si256(b[2 * i], b[2 * i + 1], 0x20));
    _mm256_store_si256((Vec256 *)(tail + 64 + 32 * i), _mm256_permute2x128_si256(b[2 * i], b[2 * i + 1], 0x31));
  }
}
// Insert into a sorted 16-byte vector with at least one trailing 255 sentinel.
inline Vec128 insert_byte(Vec128 x, U32 byte) {
  return _mm_min_epu8(x, _mm_max_epu8(_mm_slli_si128(x, 1), _mm_set1_epi8(char(byte))));
}
// Encoded counter = bucket index + 256 * count. Recover counts and flag >12 / >32.
[[gnu::always_inline]] inline U32 restore_counts(U32 *count, U32 *exceptional) {
  U32 *end = count + 256; U32 over;
  const Vec256 limit12 = _mm256_set1_epi32(12), limit32 = _mm256_set1_epi32(32);
  __asm__ volatile("vpxor %%ymm4, %%ymm4, %%ymm4\n\t"
                   ".p2align 4\n\t1:\n\t"
                   ".irp reg,0,1,2,3; vmovdqa 32*\\reg(%[count]), %%ymm\\reg; .endr\n\t"
                   ".irp reg,0,1,2,3; vpsrld $8, %%ymm\\reg, %%ymm\\reg; .endr\n\t"
                   ".irp reg,0,1,2,3; vmovdqa %%ymm\\reg, 32*\\reg(%[count]); .endr\n\t"
                   "vpmaxud %%ymm1, %%ymm0, %%ymm0; vpmaxud %%ymm3, %%ymm2, %%ymm2\n\t"
                   "vpmaxud %%ymm2, %%ymm0, %%ymm0\n\t"
                   "vpcmpgtd %[limit12], %%ymm0, %%ymm1\n\t"
                   "vmovmskps %%ymm1, %%eax; movl %%eax, (%[flags])\n\t"
                   "vpmaxud %%ymm0, %%ymm4, %%ymm4\n\t"
                   "addq $128, %[count]; addq $4, %[flags]\n\t"
                   "cmpq %[end], %[count]; jb 1b\n\t"
                   "vpcmpgtd %[limit32], %%ymm4, %%ymm0; vmovmskps %%ymm0, %k[over]\n\t"
                   : [count] "+&r"(count), [flags] "+&r"(exceptional), [over] "=&r"(over)
                   : [end] "r"(end), [limit12] "x"(limit12), [limit32] "x"(limit32)
                   : "rax", "ymm0", "ymm1", "ymm2", "ymm3", "ymm4", "cc", "memory");
  return over;
}
// All 32 counts must be <=12 and dst must have 384 writable elements.
// Overlapping 12-element stores are repaired by subsequent buckets in forward order.
[[gnu::always_inline]] inline U32 *emit32_buckets(const Byte *first, const Byte *second, const U32 *count,
                                              U32 *dst, Vec256 prefix) {
  const Vec256 step = _mm256_set1_epi32(256); const U32 *end = count + 32;
  __asm__ volatile(".p2align 4\n\t1:\n\t"
                   ".irp lane,0,1,2,3,4,5,6,7\n\t"
                   "movl 4*\\lane(%[count]), %%eax\n\t"
                   "vpmovzxbd 8*\\lane(%[first]), %%ymm1; vpmovzxbd 4*\\lane(%[second]), %%xmm2\n\t"
                   "vpor %[prefix], %%ymm1, %%ymm1; vpor %x[prefix], %%xmm2, %%xmm2\n\t"
                   "vmovdqu %%ymm1, (%[dst]); vmovdqu %%xmm2, 32(%[dst])\n\t"
                   "leaq (%[dst],%%rax,4), %[dst]\n\t"
                   "vpaddd %[step], %[prefix], %[prefix]\n\t"
                   ".endr\n\t"
                   "addq $64, %[first]; addq $32, %[second]; addq $32, %[count]\n\t"
                   "cmpq %[end], %[count]; jb 1b\n\t"
                   : [dst] "+&r"(dst), [prefix] "+&x"(prefix), [count] "+&r"(count), [first] "+&r"(first),
                     [second] "+&r"(second)
                   : [end] "r"(end), [step] "x"(step)
                   : "rax", "ymm1", "ymm2", "cc", "memory");
  return dst;
}
// end may extend past this leaf only when output cannot overwrite unread source.
inline void sort_leaf(const U16 *src, U32 n, U32 *dst, U32 base, U32 *end) {
  if (n < 64) {
    U32 copy[64];
    for (U32 i = 0; i < n; ++i) copy[i] = base | load<U16>(src + i);
    std::sort(copy, copy + n);
    std::memcpy(dst, copy, n * 4); return;
  }
  if (n > kLeafLimit) return counting_sort(src, n, dst, base);
  alignas(64) static Byte columns[kRadix * kLeafLimit];
  alignas(64) U32 count[256];
  // Only the first 12 rows need sentinels; later rows are read up to their real count.
  std::memset(columns, 255, 12 * kRadix);
  for (U32 i = 0; i < 256; ++i) count[i] = i;
  for (U32 i = 0; i < n; ++i) {
    U32 v = load<U16>(src + i), h = v >> 8, slot = count[h];
    count[h] = slot + 256; columns[slot] = Byte(v);
  }
  U32 exceptional[8];
  if (restore_counts(count, exceptional)) return counting_sort(src, n, dst, base);
  for (U32 block = 0; block < 256; block += 32) {
    alignas(32) Byte first[256], second[128];
    sort12_columns(columns + block, first, second);
    Vec256 prefix = _mm256_set1_epi32(base | (block << 8));
    if (exceptional[block >> 5] == 0 && end - dst >= 384) {
      dst = emit32_buckets(first, second, count + block, dst, prefix);
      continue;
    }
    for (U32 i = 0; i < 32; ++i) {
      U32 h = block + i, c = count[h], p = base | (h << 8);
      if (__builtin_expect(c <= 12, 1)) {
        Vec128 x = _mm_loadl_epi64((const Vec128 *)(first + 8 * i));
        U32 four = load<U32>(second + 4 * i);
        Vec256 wide = _mm256_or_si256(_mm256_cvtepu8_epi32(x), prefix);
        if (__builtin_expect(end - dst >= 12, 1)) {
          __asm__ volatile("vmovdqu {%1, %0|%0, %1}" : "=m"(*reinterpret_cast<__m256i_u *>(dst)) : "x"(wide));
          Vec128 y = _mm_or_si128(_mm_cvtepu8_epi32(_mm_cvtsi32_si128(four)), _mm256_castsi256_si128(prefix));
          _mm_storeu_si128((Vec128 *)(dst + 8), y);
        } else {
          alignas(32) U32 copy[16];
          _mm256_store_si256((Vec256 *)copy, wide);
          Vec128 y = _mm_or_si128(_mm_cvtepu8_epi32(_mm_cvtsi32_si128(four)), _mm256_castsi256_si128(prefix));
          _mm_store_si128((Vec128 *)(copy + 8), y);
          std::memcpy(dst, copy, c * 4);
        }
      } else if (c <= 16) {
        Vec128 x = _mm_loadl_epi64((const Vec128 *)(first + 8 * i));
        U32 four = load<U32>(second + 4 * i);
        x = _mm_unpacklo_epi64(x, _mm_cvtsi64_si128(uint64_t(four) | 0xffffffff00000000ull));
        for (U32 j = 12; j < c; ++j) x = insert_byte(x, columns[256 * j + h]);
        expand8(x, dst, 8, p, end);
        expand8(_mm_srli_si128(x, 8), dst + 8, c - 8, p, end);
      } else {
        Byte copy[32];
        for (U32 j = 0; j < c; ++j) copy[j] = columns[256 * j + h];
        std::sort(copy, copy + c);
        for (U32 j = 0; j < c; ++j) dst[j] = p | copy[j];
      }
      dst += c;
      prefix = _mm256_add_epi32(prefix, _mm256_set1_epi32(256));
    }
  }
}
} // namespace byte_sort

namespace fastsort {
static Byte *workspace;
inline U32 packed24_at(const Byte *p, U32 i) {
  return load<U32>(p + size_t(i / kRecordsPerBlock) * kBlockBytes + i % kRecordsPerBlock * 3) & 0xffffff;
}
inline bool buckets_fit(Byte **p, Byte **e) {
  __m256i bad = _mm256_setzero_si256();
  for (U32 k = 0; k < 256; k += 4)
    bad = _mm256_or_si256(bad, _mm256_cmpgt_epi64(_mm256_load_si256((__m256i *)(p + k)),
                                              _mm256_load_si256((__m256i *)(e + k))));
  return _mm256_testz_si256(bad, bad);
}
[[gnu::always_inline]] inline void scatter16(const Byte *record, Byte **write) {
  U32 key = record[2]; U16 lo = load<U16>(record);
  std::memcpy(write[key], &lo, 2); write[key] += 2;
}
bool split_high_byte(U32 *values, U32 n, U32 *capacity, Byte **start, U32 *count) {
  // A capacity check follows each 8192 inputs. Slack covers their worst-case spill.
  constexpr U32 guard = ((8192 + 83) / 84 + 1) * 256; // 25,344 bytes, including a final partial block.
  alignas(64) Byte cache[65536] = {};
  alignas(64) Byte *write[256], *limit[256];
  U32 cursor[256]; size_t offset = 0;
  for (U32 k = 0; k < 256; ++k) {
    start[k] = write[k] = workspace + offset;
    limit[k] = write[k] + ((size_t(capacity[k]) + 83) / 84) * 256;
    offset = size_t(limit[k] - workspace) + guard; cursor[k] = k * 256;
  }
  auto push = [&](U32 x) __attribute__((always_inline)) {
    U32 k = x >> 24, t = cursor[k];
    std::memcpy(cache + t, &x, 4); t += 3;
    if (__builtin_expect((t & 255) == 252, 0)) {
      t -= 252;
      repeat<8>([&](auto j) __attribute__((always_inline)) {
        _mm256_stream_si256((__m256i *)(write[k] + 32 * j),
                           _mm256_load_si256((const __m256i *)(cache + t + 32 * j)));
      });
      write[k] += 256;
    }
    cursor[k] = t;
  };
  for (U32 i = 0; i < n;) {
    U32 stop = std::min(n, i + 8192);
    for (; i + 8 <= stop; i += 8)
      repeat<8>([&](auto j) __attribute__((always_inline)) { push(values[i + j]); });
    for (; i < stop; ++i) push(values[i]);
    if (!buckets_fit(write, limit)) { _mm_sfence(); return false; }
  }
  for (U32 k = 0; k < 256; ++k) {
    U32 tail_bytes = cursor[k] & 255;
    count[k] = U32((write[k] - start[k]) / 256) * 84 + tail_bytes / 3;
    if (tail_bytes) { std::memcpy(write[k], cache + k * 256, 256); write[k] += 256; }
  }
  _mm_sfence();
  return buckets_fit(write, limit);
}
void sort_u16_bucket(const Byte *src, U32 n, U32 base, U32 *dst, U32 *end = nullptr) {
  byte_sort::sort_leaf((const U16 *)src, n, dst, base, end ? end : dst + n);
}
void sort_u24_bucket(const Byte *src, U32 n, U32 base, U32 *dst) {
  if (n < kSmallSort) {
    for (U32 i = 0; i < n; ++i) dst[i] = base | packed24_at(src, i);
    return std::sort(dst, dst + n);
  }
  U32 count[256] = {}, begin[257], cursor[256];
  // Exact partitioning into the output buffer; expand backwards to preserve unread input.
  if (n > 600000 || n < 16384) {
    for (U32 i = 0; i < n; ++i) ++count[packed24_at(src, i) >> 16];
    begin[0] = 0;
    for (U32 k = 0; k < 256; ++k) { cursor[k] = begin[k]; begin[k + 1] = begin[k] + count[k]; }
    for (U32 i = 0; i < n; ++i) {
      U32 x = packed24_at(src, i); U16 lo = x;
      std::memcpy((Byte *)dst + 2 * cursor[x >> 16]++, &lo, 2);
    }
    for (int k = 255; k >= 0; --k) if (count[k])
        sort_u16_bucket((Byte *)dst + 2 * begin[k], count[k], base | (U32(k) << 16), dst + begin[k]);
    return;
  }
  // Speculative capacities use odd cache-line strides; overflow triggers exact repartition.
  Byte *scratch = workspace + kScratchOffset, *write[256];
  const U32 estimated_items = (n + 255) / 256 * 5 / 4 + 32;
  const U32 stride = 64 * (((estimated_items * 2 + 63) / 64) | 1);
  U32 bytes = 0;
  for (U32 k = 0; k < 256; ++k) {
    begin[k] = bytes; write[k] = scratch + bytes;
    bytes += stride;
  }
  begin[256] = bytes;
  for (;;) {
    U32 i = 0;
    const Byte *block = src;
    for (; i + 84 <= n; i += 84, block += 256) {
      _mm_prefetch((const char *)((uintptr_t)block + 256), _MM_HINT_T0);
#pragma GCC unroll 1
      for (U32 j = 0; j < 72; j += 12)
        repeat<12>([&](auto q) __attribute__((always_inline)) {
          // One destination prefetch per three records, twelve records ahead.
          if constexpr (size_t(q) % 3 == 0) {
            U32 future = block[3 * (j + 12 + q) + 2];
            _mm_prefetch((const char *)write[future], _MM_HINT_T0);
          }
          scatter16(block + 3 * (j + q), write);
        });
      repeat<12>([&](auto q) __attribute__((always_inline)) { scatter16(block + 3 * (72 + q), write); });
    }
    for (; i < n; ++i) {
      U32 x = packed24_at(src, i); U16 lo = x;
      std::memcpy(write[x >> 16], &lo, 2); write[x >> 16] += 2;
    }
    bool bad = false;
    for (U32 k = 0; k < 256; ++k) {
      count[k] = U32(write[k] - (scratch + begin[k])) / 2;
      bad |= write[k] > scratch + begin[k + 1];
    }
    if (!bad) break;
    bytes = 0;
    for (U32 k = 0; k < 256; ++k) { begin[k] = bytes; write[k] = scratch + bytes; bytes += 2 * count[k]; }
    begin[256] = bytes;
  }
  U32 *end = dst + n;
  for (U32 k = 0; k < 256; ++k) {
    if (count[k]) sort_u16_bucket(scratch + begin[k], count[k], base | (k << 16), dst, end);
    dst += count[k];
  }
}
} // namespace fastsort

// Public judge entry point.
void sort(unsigned *values, int n) {
  using namespace fastsort;
  if (n < 2) return;
  if (n < kSmallSort || n > 100000000) return std::sort(values, values + n);
  void *raw = std::malloc(kWorkspaceBytes + 63);
  if (!raw) return std::sort(values, values + n);
  workspace = (Byte *)((uintptr_t(raw) + 63) & ~uintptr_t(63));
  U32 capacity[256], count[256];
  Byte *start[256];
  for (U32 k = 0; k < 256; ++k) capacity[k] = U32((U64(n) + 255) / 256) * 6 / 5 + 32;
  // Retry with exact capacities if the initial high-byte estimate was too small.
  if (!split_high_byte(values, n, capacity, start, count)) {
    std::memset(capacity, 0, sizeof capacity);
    for (int i = 0; i < n; ++i) ++capacity[values[i] >> 24];
    split_high_byte(values, n, capacity, start, count);
  }
  U32 offset = 0;
  for (U32 k = 0; k < 256; ++k) {
    if (count[k]) sort_u24_bucket(start[k], count[k], k << 24, values + offset);
    offset += count[k];
  }
  std::free(raw);
}

评测结果