超时求调,玄关
查看原帖
超时求调,玄关
1088663
zzzyyyyhhhhh楼主2024/4/6 19:44

TLE 70pts

#include<stdio.h>
#include<immintrin.h>
#include<iostream>
#pragma GCC target("avx2")
#define int long long

// AVX2指令集的整型变量
static unsigned mod, r, n2_;
static __m256i a0, mod1, R, hi32, ans, ml, ad;

// 扩展欧几里得算法
long long exgcd(long long a, long long b, long long &x, long long &y) {
    long long d = a;
    if (b == 0)
        x = 1, y = 0;
    else
        d = exgcd(b, a % b, y, x), y -= a / b * x;
    return d;
}

// 蒙哥马利数域内的乘法及取模
inline unsigned mul(unsigned x, unsigned y) {
    unsigned long long z = (unsigned long long)x * y;
    return (z + (unsigned long long)(unsigned(z) * r) * mod) >> 32;
}

// 将8个32位整数相加并取模
inline __m256i add(__m256i _num1, __m256i _num2) {
    // 将num1, num2内的8个整数对应相加
    __m256i apb = _mm256_add_epi32(_num1, _num2);
    // 将答案减去mod
    __m256i ret = _mm256_sub_epi32(apb, mod1);
    // 得到答案内小于0的数
    __m256i cmp = _mm256_cmpgt_epi32(a0, ret);
    __m256i add = _mm256_and_si256(cmp, mod1);
    // 将小于0的数加mod
    return _mm256_add_epi32(add, ret);
}

// 将8个32位整数相乘并取模
inline __m256i mul(__m256i _num1, __m256i _num2) {
    __m256i _num3 = _num1, _num4, _num5 = _num2;
    _num2 = _mm256_mul_epu32(_num1, _num2);
    _num1 = _mm256_mul_epu32(_mm256_mul_epu32(_num2, R), mod1);
    _num4 = _mm256_srli_epi64(_mm256_add_epi64(_num1, _num2), 32);
    _num1 = _mm256_srli_si256(_num3, 4);
    _num2 = _mm256_srli_si256(_num5, 4);
    _num2 = _mm256_mul_epu32(_num1, _num2);
    _num1 = _mm256_mul_epu32(_mm256_mul_epu32(_num2, R), mod1);
    _num1 = _mm256_and_si256(_mm256_add_epi64(_num1, _num2), hi32);
    return _mm256_or_si256(_num1, _num4);
}

// 进入蒙哥马利数域
inline unsigned mon_in(unsigned x) {
    return mul(x, n2_);
}

// 离开蒙哥马利数域
inline unsigned mon_out(unsigned x) {
    unsigned ret = ((x + (unsigned long long)(unsigned(x) * r) * mod) >> 32);
    return ret < mod ? ret : ret - mod;
}

// 计算
void solve(int n, int p) {
    unsigned i = 1;
    long long x, y;
    mod = p;
    n2_ = -(unsigned long long)mod % mod;
    exgcd(mod, 1ll << 32, x, y);
    r = -unsigned(x);

    // 预处理全0的AVX2变量(做加法有用)
    a0 = _mm256_setzero_si256();
    // 预处理全mod1,全r的AVX2变量
    mod1 = _mm256_set1_epi32(mod);
    R = _mm256_set1_epi32(r);
    // 预处理奇数位全1的AVX2变量(做乘法有用)
    hi32 = _mm256_set_epi32(-1, 0, -1, 0, -1, 0, -1, 0);
    // 初始化ans
    ans = _mm256_set1_epi32(mon_in(1));
    // 初始化全8的AVX2
    ad = _mm256_set1_epi32(mon_in(8));
    ml = _mm256_set_epi32(mon_in(8), mon_in(7), mon_in(6), mon_in(5), mon_in(4), mon_in(3), mon_in(2), mon_in(1));

    for (; i + 8 <= n; i += 8) {
        ans = mul(ans, ml);
        ml = add(ml, ad);
    }

    unsigned *fl = (unsigned *)&ans;
    unsigned as = mul(mul(mul(fl[0], fl[1]), mul(fl[2], fl[3])), mul(mul(fl[4], fl[5]), mul(fl[6], fl[7])));
    as = mon_out(as);

    for (; i <= n; i++)
        as = 1ull * as * i % mod;

    if (as)
        std::cout << "No";
    else
        std::cout << "Yes";
}

signed main() {
    int n, p = 12;
    scanf("%lld", &n);
    solve(n, p);
    return 0;
}
2024/4/6 19:44
加载中...