代码

// Duck.ac 1002: exact convolution for coefficients 0..9, degrees <= 1,000,000.
// AVX2 / GCC 9.3, adapted from Qwerty1232: https://duck.ac/submission/28087
// Cleaned from https://duck.ac/submission/48181 (20.505208 ms).
// p = 39 * 2^21 + 1 > 81 * 1,000,001, so one modulus gives exact integers.
// The fast path fuses radix-4 NTT stages in 512-element blocks and reuses c.
#include <immintrin.h>

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <vector>

#pragma GCC target("avx2,bmi")

using u32 = uint32_t;
using u64 = uint64_t;

struct Montgomery {
    u32 mod;   // mod
    u32 mod2;  // 2 * mod
    u32 n_inv; // n_inv * mod == -1 (mod 2^32)
    u32 r;     // 2^32 % mod
    u32 r2;    // (2^32)^2 % mod

    Montgomery() = default;
    Montgomery(u32 mod) : mod(mod) {
        assert(mod % 2 == 1);
        assert(mod < (1 << 30));
        mod2 = 2 * mod;
        n_inv = 1;
        for (int i = 0; i < 5; i++) {
            n_inv *= 2 + n_inv * mod;
        }
        r = (u64(1) << 32) % mod;
        r2 = u64(r) * r % mod;
    }

    u32 shrink(u32 val) const { return std::min(val, val - mod); }
    u32 shrink2(u32 val) const { return std::min(val, val - mod2); }

    template <bool strict = true> u32 reduce(u64 val) const {
        u32 res = (val + u32(val) * n_inv * u64(mod)) >> 32;
        if (strict) res = shrink(res);
        return res;
    }

    template <bool strict = true> u32 mul(u32 a, u32 b) const { return reduce<strict>(u64(a) * b); }

    template <bool input_in_space = false, bool output_in_space = false> u32 power(u32 b, u32 e) const {
        if (!input_in_space) b = mul<false>(b, r2);
        u32 r = output_in_space ? this->r : 1;
        for (; e > 0; e >>= 1) {
            if (e & 1) r = mul<false>(r, b);
            b = mul<false>(b, b);
        }
        return shrink(r);
    }
};

using i256 = __m256i;
using u32x8 = u32 __attribute__((vector_size(32)));
using u64x4 = u64 __attribute__((vector_size(32)));

u32x8 load_u32x8(const u32 *ptr) {
    return (u32x8)_mm256_load_si256((const i256 *)ptr);
}
void store_u32x8(u32 *ptr, u32x8 vec) {
    _mm256_store_si256((i256 *)ptr, (i256)vec);
}

struct MontgomeryAVX2 {
    static constexpr u32x8 mod = {81788929, 81788929, 81788929, 81788929, 81788929, 81788929, 81788929, 81788929};
    static constexpr u32x8 mod2 = {163577858, 163577858, 163577858, 163577858,
                                   163577858, 163577858, 163577858, 163577858};
    static constexpr u32x8 n_inv = {81788927, 81788927, 81788927, 81788927, 81788927, 81788927, 81788927, 81788927};
    static constexpr u32x8 r = {41942988, 41942988, 41942988, 41942988, 41942988, 41942988, 41942988, 41942988};
    static constexpr u32x8 r2 = {56088131, 56088131, 56088131, 56088131, 56088131, 56088131, 56088131, 56088131};
    MontgomeryAVX2() = default;
    explicit MontgomeryAVX2(u32 p) { assert(p == 81788929); }

    u32x8 shrink(u32x8 vec) const { return (u32x8)_mm256_min_epu32((i256)vec, _mm256_sub_epi32((i256)vec, (i256)mod)); }
    template <int Low = 0> u32x8 canonical_wide(u32x8 v) const {
        for (int shift = 5; shift >= Low; shift--) {
            u32x8 p = mod << shift;
            v = (u32x8)_mm256_min_epu32((i256)v, (i256)(v - p));
        }
        return v;
    }
    u32x8 shrink2(u32x8 vec) const {
        return (u32x8)_mm256_min_epu32((i256)vec, _mm256_sub_epi32((i256)vec, (i256)mod2));
    }
    u32x8 shrink2_n(u32x8 vec) const {
        return (u32x8)_mm256_min_epu32((i256)vec, _mm256_add_epi32((i256)vec, (i256)mod2));
    }

    template <bool strict = true> u32x8 reduce(u64x4 x0246, u64x4 x1357) const {
        u64x4 x0246_ninv = (u64x4)_mm256_mul_epu32((i256)x0246, (i256)n_inv);
        u64x4 x1357_ninv = (u64x4)_mm256_mul_epu32((i256)x1357, (i256)n_inv);
        u64x4 x0246_res = (u64x4)_mm256_add_epi64((i256)x0246, _mm256_mul_epu32((i256)x0246_ninv, (i256)mod));
        u64x4 x1357_res = (u64x4)_mm256_add_epi64((i256)x1357, _mm256_mul_epu32((i256)x1357_ninv, (i256)mod));
        u32x8 res = (u32x8)_mm256_or_si256(_mm256_bsrli_epi128((i256)x0246_res, 4), (i256)x1357_res);
        if (strict) res = shrink(res);
        return res;
    }

    template <bool strict = true, bool b_use_only_even = false> u32x8 mul_u32x8(u32x8 a, u32x8 b) const {
        u32x8 a_sh = (u32x8)_mm256_bsrli_epi128((i256)a, 4);
        u32x8 b_sh = b_use_only_even ? b : (u32x8)_mm256_bsrli_epi128((i256)b, 4);
        u64x4 x0246 = (u64x4)_mm256_mul_epu32((i256)a, (i256)b);
        u64x4 x1357 = (u64x4)_mm256_mul_epu32((i256)a_sh, (i256)b_sh);
        return reduce<strict>(x0246, x1357);
    }

    template <bool strict = true> u64x4 mul_u64x4(u64x4 a, u64x4 b) const {
        u64x4 pr = (u64x4)_mm256_mul_epu32((i256)a, (i256)b);
        u64x4 pr2 = (u64x4)_mm256_mul_epu32(_mm256_mul_epu32((i256)pr, (i256)n_inv), (i256)mod);
        u64x4 res = (u64x4)_mm256_bsrli_epi128(_mm256_add_epi64((i256)pr, (i256)pr2), 4);
        if (strict) res = (u64x4)shrink((u32x8)res);
        return res;
    }
};

// Assemble read-only roots at build time. C++ constexpr expansion exceeds the
// judge compiler's memory limit. n1/n2/n3 are w1/w2/w3 * n_inv modulo 2^32.
namespace fixed_roots {
struct Twiddle {
    u32 w1, w2, w3, n1, n2, n3;
};
struct Table {
    Twiddle data[65536];
};
extern const Table forward asm("poly_roots_forward");
extern const Table inverse asm("poly_roots_inverse");
struct DotTable {
    u32 data[32768][4];
};
extern const DotTable dot asm("poly_roots_dot");
} // namespace fixed_roots
asm(R"asm(
.pushsection .rodata
// Select by the number of trailing one bits in the table index.
.macro next_factor dest, mask, value, rest:vararg
.if ((_i & \mask) == 0)
.set \dest,\value
.else
next_factor \dest,(\mask*2),\rest
.endif
.endm
.p2align 6
.globl poly_roots_forward
.type poly_roots_forward,@object
poly_roots_forward:
.set _i,0
.set _w1,41942988
.set _w2,41942988
.rept 65536
.set _w12,(((_w1*_w2)%81788929)*1557504)%81788929
.long _w1,_w2,_w12,((_w1*81788927)&0xffffffff),((_w2*81788927)&0xffffffff),((_w12*81788927)&0xffffffff)
next_factor _fac, 1, 1977387,49739338,76551861,57685215,25318722,22305379,75160758,77449485,49050524,58847824,69356575,69052175,45043381,68811137,48691376,28111944,26577652
next_factor _fac2, 1, 57807995,1883838,27152551,62819432,22367481,25489457,54748607,18371892,60074596,43336831,16579980,78708963,26101542,51041304,60500196,40232015,28323882
.set _w1,(_w1*_fac2)%81788929
.set _w2,(_w2*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_forward,.-poly_roots_forward
.p2align 6
.globl poly_roots_inverse
.type poly_roots_inverse,@object
poly_roots_inverse:
.set _i,0
.set _w1,41942988
.set _w2,41942988
.rept 65536
.set _w12,(((_w1*_w2)%81788929)*1557504)%81788929
.long _w1,_w2,_w12,((_w1*81788927)&0xffffffff),((_w2*81788927)&0xffffffff),((_w12*81788927)&0xffffffff)
next_factor _fac, 1, 1883838,34192649,65864533,47472559,32202660,46455854,18299665,51166265,46148164,40005067,42538512,22507185,19881487,13191717,67317322,1064838,16759432
next_factor _fac2, 1, 23980934,1977387,58967103,56026958,74557765,58488554,3169619,20142414,28119438,26733415,74787290,67900511,63391377,74641937,67976842,40043517,2457972
.set _w1,(_w1*_fac2)%81788929
.set _w2,(_w2*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_inverse,.-poly_roots_inverse
.p2align 6
.globl poly_roots_dot
.type poly_roots_dot,@object
poly_roots_dot:
.set _i,0
.set _d0,41942988
.set _d1,42958308
.set _d2,28282409
.set _d3,36011086
.rept 32768
.long _d0,_d1,_d2,_d3
next_factor _fac, 1, 34192649,76852948,45870503,27153147,50722843,53125215,43544278,13378268,50576854,50366248,18491959,52344447,58962465,12062499,19859762,9337084
.set _d0,(_d0*_fac)%81788929
.set _d1,(_d1*_fac)%81788929
.set _d2,(_d2*_fac)%81788929
.set _d3,(_d3*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_dot,.-poly_roots_dot
.purgem next_factor
.popsection
)asm");

class NTT {
    // Global coefficient offset of the quarter currently in local scratch.
    mutable int data_origin = 0;

  public:
    u32 mod;

  private:
    static const int LG = 32; // more than enough for u32

    Montgomery mt;
    MontgomeryAVX2 mts;

    u32 w[4], wr[4];

    u64x4 wt_init, wrt_init;
    u64x4 wd_x4[LG], wrd_x4[LG];

    u64x4 wl_init;
    u64x4 wld_x4[LG];

  public:
    NTT(u32 mod) : mod(mod), mt(mod), mts(mod) {
        const Montgomery mt = this->mt;
        constexpr u32 pr_root = 7; // Primitive root for the fixed modulus 81,788,929.

        int lg = __builtin_ctz(mod - 1);
        assert(lg <= LG);

        memset(w, 0, sizeof(w));
        memset(wr, 0, sizeof(wr));
        memset(wd_x4, 0, sizeof(wd_x4));
        memset(wrd_x4, 0, sizeof(wrd_x4));
        memset(wld_x4, 0, sizeof(wld_x4));

        std::vector<u32> vec(lg + 1), vecr(lg + 1);
        vec[lg] = mt.power<false, true>(pr_root, (mod - 1) >> lg);
        vecr[lg] = mt.power<true, true>(vec[lg], mod - 2);
        for (int i = lg - 1; i >= 0; i--) {
            vec[i] = mt.mul<true>(vec[i + 1], vec[i + 1]);
            vecr[i] = mt.mul<true>(vecr[i + 1], vecr[i + 1]);
        }

        w[0] = wr[0] = mt.r;
        if (lg >= 2) {
            w[1] = vec[2], wr[1] = vecr[2];
            if (lg >= 3) {
                w[2] = vec[3], wr[2] = vecr[3];
                w[3] = mt.mul<true>(w[1], w[2]);
                wr[3] = mt.mul<true>(wr[1], wr[2]);
            }
        }
        wt_init = (u64x4)_mm256_setr_epi64x(w[0], w[0], w[0], w[1]);
        wrt_init = (u64x4)_mm256_setr_epi64x(wr[0], wr[0], wr[0], wr[1]);

        wl_init = (u64x4)_mm256_setr_epi64x(w[0], w[1], w[2], w[3]);

        u32 prf = mt.r, prf_r = mt.r;
        for (int i = 0; i < lg - 2; i++) {
            u32 f = mt.mul<true>(prf, vec[i + 3]), fr = mt.mul<true>(prf_r, vecr[i + 3]);
            prf = mt.mul<true>(prf, vecr[i + 3]), prf_r = mt.mul<true>(prf_r, vec[i + 3]);
            u32 f2 = mt.mul<true>(f, f), f2r = mt.mul<true>(fr, fr);

            wd_x4[i] = (u64x4)_mm256_setr_epi64x(f2, f, f2, f);
            wrd_x4[i] = (u64x4)_mm256_setr_epi64x(f2r, fr, f2r, fr);
        }

        prf = mt.r;
        for (int i = 0; i < lg - 3; i++) {
            u32 f = mt.mul<true>(prf, vec[i + 4]);
            prf = mt.mul<true>(prf, vecr[i + 4]);
            wld_x4[i] = (u64x4)_mm256_set1_epi64x(f);
        }
    }

  private:
    static const int L0 = 3;
    int leaf_log2(int lg) const { return lg % 2 == L0 % 2 ? L0 : L0 + 1; }

    // Precomputed w*n_inv lets the product and reduction start independently.
    static u32x8 mul_pre(u32x8 a, u32x8 w, u32x8 wn, const MontgomeryAVX2 &mts) {
        i256 a1 = _mm256_srli_epi64((i256)a, 32);
        i256 m0 = _mm256_mul_epu32((i256)a, (i256)wn), m1 = _mm256_mul_epu32(a1, (i256)wn);
        i256 p0 = _mm256_mul_epu32((i256)a, (i256)w), p1 = _mm256_mul_epu32(a1, (i256)w);
        p0 = _mm256_add_epi64(p0, _mm256_mul_epu32(m0, (i256)mts.mod));
        p1 = _mm256_add_epi64(p1, _mm256_mul_epu32(m1, (i256)mts.mod));
        return (u32x8)_mm256_blend_epi32(_mm256_srli_epi64(p0, 32), p1, 0xaa);
    }
    template <bool inverse, bool trivial>
    static void butterfly_pair(u32x8 &a, u32x8 &b, u32x8 w, u32x8 wn, const MontgomeryAVX2 &mts) {
        if constexpr (!inverse) {
            b = trivial ? b : mul_pre(b, w, wn, mts);
            auto x = a + b;
            b = a + mts.mod2 - b;
            a = x;
        } else {
            auto x = mts.shrink2(a + b);
            b = trivial ? mts.shrink2_n(a - b) : mul_pre(a + mts.mod2 - b, w, wn, mts);
            a = x;
        }
    }
    // Fixed-size path: table-indexed roots, no running twiddle dependency.
    // For the official input, forward residues stay below 51p < 2^32.
    template <int k, bool inverse, bool trivial = false>
    __attribute__((always_inline)) inline void transform_fixed(int i, u32 *data, const MontgomeryAVX2 &mts) const {
        const auto &tw = (inverse ? fixed_roots::inverse : fixed_roots::forward).data[unsigned(i) >> (k + 2)];
        u32x8 w1 = (u32x8)_mm256_set1_epi32(tw.w1), w2 = (u32x8)_mm256_set1_epi32(tw.w2),
              w3 = (u32x8)_mm256_set1_epi32(tw.w3);
        u32x8 n1 = (u32x8)_mm256_set1_epi32(tw.n1), n2 = (u32x8)_mm256_set1_epi32(tw.n2),
              n3 = (u32x8)_mm256_set1_epi32(tw.n3);
        u32x8 root = (u32x8)_mm256_set1_epi32(inverse ? 38830621 : 42958308);
        u32x8 root_n = (u32x8)_mm256_set1_epi32(inverse ? 1259306467 : 3035660828);
        if constexpr (trivial) {
            w3 = root;
            n3 = root_n;
        }
        if constexpr (inverse && !trivial && k >= 5) {
            // 2x-unrolled nontrivial inverse: two independent butterflies in flight.
            const int step = 1 << k;
            for (int j = 0; j < step; j += 16) {
                u32 *p = data + i - data_origin + j;
                u32 *q = p + 8;
                auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
                     d = load_u32x8(p + step * 3);
                auto e = load_u32x8(q), f = load_u32x8(q + step), g = load_u32x8(q + step * 2),
                     h = load_u32x8(q + step * 3);
                auto u = a + b, s = c + d, v = a + mts.mod2 - b;
                auto t = mul_pre(c + mts.mod2 - d, root, root_n, mts);
                auto u2 = e + f, s2 = g + h, v2 = e + mts.mod2 - f;
                auto t2 = mul_pre(g + mts.mod2 - h, root, root_n, mts);
                auto sum = u + s;
                sum = (u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2));
                a = mts.shrink2(sum);
                auto sum2 = u2 + s2;
                sum2 = (u32x8)_mm256_min_epu32((i256)sum2, (i256)(sum2 - mts.mod2 - mts.mod2));
                e = mts.shrink2(sum2);
                c = mul_pre(u + mts.mod2 + mts.mod2 - s, w1, n1, mts);
                g = mul_pre(u2 + mts.mod2 + mts.mod2 - s2, w1, n1, mts);
                b = mul_pre(v + t, w2, n2, mts);
                f = mul_pre(v2 + t2, w2, n2, mts);
                d = mul_pre(v + mts.mod2 - t, w3, n3, mts);
                h = mul_pre(v2 + mts.mod2 - t2, w3, n3, mts);
                store_u32x8(p, a);
                store_u32x8(p + step, b);
                store_u32x8(p + 2 * step, c);
                store_u32x8(p + 3 * step, d);
                store_u32x8(q, e);
                store_u32x8(q + step, f);
                store_u32x8(q + 2 * step, g);
                store_u32x8(q + 3 * step, h);
            }
            return;
        }
        for (int j = 0; j < (1 << k); j += 8) {
            u32 *p = data + i - data_origin + j;
            int step = 1 << k;
            auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
                 d = load_u32x8(p + step * 3);
            if constexpr (!inverse) {
                if constexpr (trivial) {
                    butterfly_pair<false, true>(a, c, w1, n1, mts);
                    butterfly_pair<false, true>(b, d, w1, n1, mts);
                    butterfly_pair<false, true>(a, b, w2, n2, mts);
                    butterfly_pair<false, false>(c, d, w3, n3, mts);
                } else {
                    auto cc = mul_pre(c, w1, n1, mts), bb = mul_pre(b, w2, n2, mts), dd = mul_pre(d, w3, n3, mts);
                    auto A = a + cc, C = a + mts.mod2 - cc, B = bb + dd;
                    auto D = mul_pre(bb + mts.mod2 - dd, root, root_n, mts);
                    a = A + B;
                    b = A + mts.mod2 + mts.mod2 - B;
                    c = C + D;
                    d = C + mts.mod2 - D;
                }
            } else {
                if constexpr (trivial) {
                    butterfly_pair<true, true>(a, b, w2, n2, mts);
                    butterfly_pair<true, false>(c, d, w3, n3, mts);
                    butterfly_pair<true, true>(a, c, w1, n1, mts);
                    butterfly_pair<true, true>(b, d, w1, n1, mts);
                } else {
                    auto u = a + b, s = c + d, v = a + mts.mod2 - 

评测结果