代码

// 1001a: sort n unsigned in place (interface: void sort(unsigned *a, int n)).
//
// Strategy: MSD partition by the top 10 bits into 1024 small buckets (~10 elements
// each for random 32-bit input), then sort every bucket with an AVX2 bitonic sorting
// network (16 lanes, padded with 0xFFFFFFFF).  A bucket larger than 16 elements is
// finished with std::sort.
//
// Why this beats a radix sort here: 4-pass LSD radix needs 12 stores per element
// (4 histogram + 4x(1 counter + 1 data)); the scatter passes alone cost ~33k ticks
// each on the judge.  This version needs ~1 histogram store, 1 scatter counter store,
// 1 scatter data store, and 2 in-bucket stores per element, and the in-bucket work is
// register-resident.  Judge-HW measurement (duck.ac custom-test channel, best of 300,
// 10000 random u32): 4-pass radix 188k ticks -> this 155k ticks.
//
// A 4-pass LSD radix is kept as a fallback for CPUs without AVX2, and small inputs
// use insertion sort.
#include <algorithm>
#include <cstring>
#include <immintrin.h>

typedef unsigned u32;
typedef unsigned short u16;

#if defined(__GNUC__)
#define AVX2FN __attribute__((target("avx2")))
#else
#define AVX2FN
#endif

#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC push_options
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#endif

// ---------------- AVX2 primitives ----------------
static const int IDX1[8] = {1,0,3,2,5,4,7,6};
static const int IDX2[8] = {2,3,0,1,6,7,4,5};
static const int IDX4[8] = {4,5,6,7,0,1,2,3};
static const int IDXR[8] = {7,6,5,4,3,2,1,0};

#define SEL_A _mm256_setr_epi32(0,-1,-1,0, 0,-1,-1,0)   // bitonic stage j=1,k=2
#define SEL_B _mm256_setr_epi32(0,0,-1,-1, -1,-1,0,0)   // j=2,k=4
#define SEL_C _mm256_setr_epi32(0,-1,0,-1, -1,0,-1,0)   // j=1,k=4
#define SEL_D _mm256_setr_epi32(0,0,0,0, -1,-1,-1,-1)   // j=4,k=8
#define SEL_E _mm256_setr_epi32(0,0,-1,-1, 0,0,-1,-1)   // j=2,k=8
#define SEL_F _mm256_setr_epi32(0,-1,0,-1, 0,-1,0,-1)   // j=1,k=8

AVX2FN static inline __m256i ce_stage(__m256i v, const int* idx, __m256i sel) {
    __m256i s = _mm256_permutevar8x32_epi32(v, _mm256_loadu_si256((const __m256i*)(const void*)idx));
    __m256i mn = _mm256_min_epu32(v, s), mx = _mm256_max_epu32(v, s);
    return _mm256_blendv_epi8(mn, mx, sel);
}
AVX2FN static inline __m256i net8(__m256i v) {
    v = ce_stage(v, IDX1, SEL_A);
    v = ce_stage(v, IDX2, SEL_B);
    v = ce_stage(v, IDX1, SEL_C);
    v = ce_stage(v, IDX4, SEL_D);
    v = ce_stage(v, IDX2, SEL_E);
    v = ce_stage(v, IDX1, SEL_F);
    return v;
}
AVX2FN static inline void merge16(__m256i& lo, __m256i& hi) {
    __m256i hr = _mm256_permutevar8x32_epi32(hi, _mm256_loadu_si256((const __m256i*)(const void*)IDXR));
    __m256i mn = _mm256_min_epu32(lo, hr), mx = _mm256_max_epu32(lo, hr);
    lo = ce_stage(ce_stage(ce_stage(mn, IDX4, SEL_D), IDX2, SEL_E), IDX1, SEL_F);
    hi = ce_stage(ce_stage(ce_stage(mx, IDX4, SEL_D), IDX2, SEL_E), IDX1, SEL_F);
}
AVX2FN static inline __m256i load_pad8(const u32* p, int cnt) {
    __m256i v = _mm256_loadu_si256((const __m256i*)(const void*)p);
    __m256i idx = _mm256_setr_epi32(0,1,2,3,4,5,6,7);
    __m256i m = _mm256_cmpgt_epi32(_mm256_set1_epi32(cnt), idx);
    return _mm256_blendv_epi8(_mm256_set1_epi32(-1), v, m);
}
AVX2FN static inline void store_part(u32* p, __m256i v, int cnt) {
    __m256i idx = _mm256_setr_epi32(0,1,2,3,4,5,6,7);
    __m256i m = _mm256_cmpgt_epi32(_mm256_set1_epi32(cnt), idx);
    _mm256_maskstore_epi32((int*)(void*)p, m, v);
}
// sorts a bucket of size 1..16 in place
AVX2FN static void sort16_pad(u32* p, int size) {
    int c0 = size < 8 ? size : 8;
    __m256i lo = net8(load_pad8(p, c0));
    if (size <= 8) { store_part(p, lo, c0); return; }
    int c1 = size - 8;
    __m256i hi = net8(load_pad8(p + 8, c1));
    merge16(lo, hi);
    store_part(p, lo, c0);
    store_part(p + 8, hi, c1);
}

#if defined(__GNUC__) && !defined(__clang__)
#pragma GCC pop_options
#endif

// ---------------- fallbacks ----------------
static void radix_sort(u32* a, int n) {
    static u16 h[4][256];
    static u32 tmp[1 << 16];
    for (int i = 0; i < 256; i++) { h[0][i] = 0; h[1][i] = 0; h[2][i] = 0; h[3][i] = 0; }
    for (int i = 0; i < n; i++) {
        u32 v = a[i];
        h[0][v & 255]++; h[1][(v >> 8) & 255]++; h[2][(v >> 16) & 255]++; h[3][v >> 24]++;
    }
    u32 *src = a, *dst = tmp;
    int done = 0;
    for (int p = 0; p < 4; p++) {
        int nb = 0;
        for (int k = 0; k < 256; k++) if (h[p][k]) { nb = 1; break; }
        if (!nb) continue;
        unsigned s = 0;
        u16* hp = h[p];
        for (int k = 0; k < 256; k++) { unsigned c = hp[k]; hp[k] = (u16)s; s += c; }
        int sh = p * 8;
        for (int i = 0; i < n; i++) { u32 v = src[i]; dst[hp[(v >> sh) & 255]++] = v; }
        u32* t = src; src = dst; dst = t;
        done++;
    }
    if (done && src != a) for (int i = 0; i < n; i++) a[i] = src[i];
}

// ---------------- main entry ----------------
#define NBUCK 1024
static u16 g_cnt[NBUCK];
static u32 g_tmp[1 << 17];

// returns 1 on success, 0 if the input is too skewed for the bucket path
AVX2FN static int bucket_sort(u32* a, int n) {
    for (int k = 0; k < NBUCK; k++) g_cnt[k] = 0;
    for (int i = 0; i < n; i++) g_cnt[a[i] >> 22]++;
    unsigned s = 0;
    int nonEmpty = 0, maxc = 0;
    for (int k = 0; k < NBUCK; k++) {
        unsigned c = g_cnt[k];
        if (c) { nonEmpty++; if ((int)c > maxc) maxc = (int)c; }
        g_cnt[k] = (u16)s; s += c;
    }
    if (maxc > 32 || nonEmpty < 256) return 0;   // too skewed: caller uses radix
    for (int i = 0; i < n; i++) { u32 v = a[i]; g_tmp[g_cnt[v >> 22]++] = v; }
    unsigned start = 0;
    for (int k = 0; k < NBUCK; k++) {
        int c = (int)(unsigned)(g_cnt[k] - start);
        if (c > 0) {
            if (c <= 16) sort16_pad(g_tmp + start, c);
            else std::sort(g_tmp + start, g_tmp + start + c);
        }
        start = g_cnt[k];
    }
    memcpy(a, g_tmp, (size_t)n * 4);
    return 1;
}

void sort(unsigned* a, int n) {
    if (n <= 1) return;
    if (n <= 24) {                      // tiny: branchless-free simple insertion sort
        for (int i = 1; i < n; i++) {
            u32 v = a[i]; int j = i - 1;
            while (j >= 0 && a[j] > v) { a[j + 1] = a[j]; j--; }
            a[j + 1] = v;
        }
        return;
    }
    if (n <= (1 << 17) - 16 && __builtin_cpu_supports("avx2") && bucket_sort(a, n)) return;
    radix_sort(a, n);
}

评测结果