代码

/* 【WC2017】挑战-任务2: rock-paper-scissors between two rows.
 * Query (x,y,l): count i<l with s1[x+i] beating s2[y+i], where 0=rock,1=scissors,2=paper
 * and a beats b iff b == (a+1)%3  ->  winning character pairs (0,1),(1,2),(2,0).
 *
 * Bitset formulation: A[c][k] = 1 iff s1[64k+bit] == c, likewise B[c].  Then
 *   answer = popcount( (A0>>x & B1>>y) | (A1>>x & B2>>y) | (A2>>x & B0>>y) )   over l bits
 * (the three terms are disjoint, so OR == sum).  Queries are bucketed by
 * (x&63, y&63); inside a bucket the six arrays are shifted once, after which every
 * query is a straight AVX2 overlap-popcount with aligned pointers. */
#include <immintrin.h>
#include <stdlib.h>
#include <string.h>

typedef unsigned long long u64;
typedef unsigned u32;

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

static inline __m256i pc256(__m256i v) {
    const __m256i lut = _mm256_setr_epi8(0,1,1,2,1,2,2,3,1,2,2,3,2,3,3,4,
                                         0,1,1,2,1,2,2,3,1,2,2,3,2,3,3,4);
    const __m256i m = _mm256_set1_epi8(0x0f);
    __m256i lo = _mm256_and_si256(v, m);
    __m256i hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), m);
    return _mm256_add_epi8(_mm256_shuffle_epi8(lut, lo), _mm256_shuffle_epi8(lut, hi));
}

#define WORDS(n) (((n) + 63) / 64 + 3)

void solve(int n, int q, char *s1, char *s2, int *q_x, int *q_y, int *q_len, unsigned *ans) {
    if (q <= 0) return;
    if (n <= 4096) {                       /* tiny case: straight scan */
        for (int t = 0; t < q; t++) {
            int x = q_x[t], y = q_y[t], l = q_len[t];
            u32 c = 0;
            for (int i = 0; i < l; i++) {
                int a = s1[x + i] - '0', b = s2[y + i] - '0';
                c += (b == (a + 1) % 3);
            }
            ans[t] = c;
        }
        return;
    }
    int W = WORDS(n);
    u64 *A[3], *B[3];
    for (int c = 0; c < 3; c++) {
        A[c] = (u64 *)calloc((size_t)W, 8);
        B[c] = (u64 *)calloc((size_t)W, 8);
    }
    if (!A[0] || !A[1] || !A[2] || !B[0] || !B[1] || !B[2]) {
        for (int t = 0; t < q; t++) ans[t] = 0;
        return;
    }
    for (int i = 0; i < n; i++) { int a = s1[i] - '0'; if ((unsigned)a < 3u) A[a][i >> 6] |= 1ull << (i & 63); }
    for (int i = 0; i < n; i++) { int b = s2[i] - '0'; if ((unsigned)b < 3u) B[b][i >> 6] |= 1ull << (i & 63); }

    /* bucket queries by (x&63)*64+(y&63) with a counting sort */
    int cntb[4096 + 1];
    memset(cntb, 0, sizeof(cntb));
    for (int t = 0; t < q; t++) cntb[((q_x[t] & 63) << 6) | (q_y[t] & 63)]++;
    int off[4097];
    int s = 0;
    for (int k = 0; k < 4096; k++) { off[k] = s; s += cntb[k]; }
    off[4096] = s;
    int *ord = (int *)malloc((size_t)q * sizeof(int));
    if (!ord) { for (int t = 0; t < q; t++) ans[t] = 0; return; }
    {
        int pos[4096];
        memcpy(pos, off, sizeof(pos));
        for (int t = 0; t < q; t++) ord[pos[((q_x[t] & 63) << 6) | (q_y[t] & 63)]++] = t;
    }

    u64 *SA[3], *SB[3];
    for (int c = 0; c < 3; c++) {
        SA[c] = (u64 *)_mm_malloc((size_t)W * 8, 32);
        SB[c] = (u64 *)_mm_malloc((size_t)W * 8, 32);
    }
    const __m256i zero = _mm256_setzero_si256();

    for (int g = 0; g < 4096; g++) {
        int b0 = off[g], b1 = off[g + 1];
        if (b0 >= b1) continue;
        int rx = g >> 6, ry = g & 63;
        /* shift the six arrays right by rx / ry bits (window read at word offsets) */
        for (int c = 0; c < 3; c++) {
            const u64 *a = A[c]; u64 *d = SA[c];
            const u64 *b = B[c]; u64 *e = SB[c];
            if (rx == 0) memcpy(d, a, (size_t)W * 8);
            else {
                __m256i sh = _mm256_set1_epi64x(rx), iv = _mm256_set1_epi64x(64 - rx);
                int k = 0;
                for (; k + 4 <= W - 1; k += 4) {
                    __m256i x0 = _mm256_loadu_si256((const __m256i *)(const void *)(a + k));
                    __m256i x1 = _mm256_loadu_si256((const __m256i *)(const void *)(a + k + 1));
                    _mm256_storeu_si256((__m256i *)(void *)(d + k),
                                        _mm256_or_si256(_mm256_srlv_epi64(x0, sh), _mm256_sllv_epi64(x1, iv)));
                }
                for (; k < W - 1; k++) d[k] = (a[k] >> rx) | (a[k + 1] << (64 - rx));
            }
            if (ry == 0) memcpy(e, b, (size_t)W * 8);
            else {
                __m256i sh = _mm256_set1_epi64x(ry), iv = _mm256_set1_epi64x(64 - ry);
                int k = 0;
                for (; k + 4 <= W - 1; k += 4) {
                    __m256i x0 = _mm256_loadu_si256((const __m256i *)(const void *)(b + k));
                    __m256i x1 = _mm256_loadu_si256((const __m256i *)(const void *)(b + k + 1));
                    _mm256_storeu_si256((__m256i *)(void *)(e + k),
                                        _mm256_or_si256(_mm256_srlv_epi64(x0, sh), _mm256_sllv_epi64(x1, iv)));
                }
                for (; k < W - 1; k++) e[k] = (b[k] >> ry) | (b[k + 1] << (64 - ry));
            }
        }
        for (int t = b0; t < b1; t++) {
            int qi = ord[t];
            int x = q_x[qi], y = q_y[qi], l = q_len[qi];
            const u64 *p0 = SA[0] + (x >> 6), *p1 = SA[1] + (x >> 6), *p2 = SA[2] + (x >> 6);
            const u64 *r0 = SB[0] + (y >> 6), *r1 = SB[1] + (y >> 6), *r2 = SB[2] + (y >> 6);
            int nw = l >> 6, rem = l & 63;
            __m256i acc = zero;
            int j = 0;
            for (; j + 4 <= nw; j += 4) {
                __m256i a0 = _mm256_loadu_si256((const __m256i *)(const void *)(p0 + j));
                __m256i a1 = _mm256_loadu_si256((const __m256i *)(const void *)(p1 + j));
                __m256i a2 = _mm256_loadu_si256((const __m256i *)(const void *)(p2 + j));
                __m256i b0v = _mm256_loadu_si256((const __m256i *)(const void *)(r0 + j));
                __m256i b1v = _mm256_loadu_si256((const __m256i *)(const void *)(r1 + j));
                __m256i b2v = _mm256_loadu_si256((const __m256i *)(const void *)(r2 + j));
                __m256i m = _mm256_or_si256(
                    _mm256_or_si256(_mm256_and_si256(a0, b1v), _mm256_and_si256(a1, b2v)),
                    _mm256_and_si256(a2, b0v));
                acc = _mm256_add_epi64(acc, _mm256_sad_epu8(pc256(m), zero));
            }
            u32 c = 0;
            {
                u64 tmp[4];
                _mm256_storeu_si256((__m256i *)(void *)tmp, acc);
                c = (u32)(tmp[0] + tmp[1] + tmp[2] + tmp[3]);
            }
            for (; j < nw; j++) {
                u64 m = (p0[j] & r1[j]) | (p1[j] & r2[j]) | (p2[j] & r0[j]);
                c += (u32)__builtin_popcountll(m);
            }
            if (rem) {
                u64 m = ((p0[nw] & r1[nw]) | (p1[nw] & r2[nw]) | (p2[nw] & r0[nw])) & ((1ull << rem) - 1);
                c += (u32)__builtin_popcountll(m);
            }
            ans[qi] = c;
        }
    }
    for (int c = 0; c < 3; c++) { _mm_free(SA[c]); _mm_free(SB[c]); free(A[c]); free(B[c]); }
    free(ord);
}

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

评测结果