【未完善】多项式全家桶

发布时间 2023-11-26 16:51:12作者: 383494
#include <iostream>
#include <cmath>
#include <cctype>
#include <functional>
#include <algorithm>
#include <vector>
#define UP(i,s,e) for(auto i=s; i<e; ++i)
#define DOWN(i,e,s) for(auto i=e; i-->s; )
using std::cin; using std::cout;
using u32 = unsigned int;
using u64 = unsigned long long;
constexpr u32 MOD = 998244353, G = 3, TIM = 23;
template<typename T, typename U>
constexpr T qpow(T x, U tim){
    T ans = 1;
    while(tim){
        if(tim & 1) ans *= x;
        x *= x;
        tim >>= 1;
    }
    return ans;
}
struct m32{
    u32 val;
    m32(){}
    constexpr m32(u32 x):val(x){}
    constexpr bool operator==(m32 b){ return val == b.val; }
    constexpr m32 operator+=(m32 b){ val += b.val; val %= MOD; return *this;}
    constexpr m32 operator*=(m32 b){ val = (u64)val * b.val % MOD; return *this;}
    constexpr m32 operator-=(m32 b){ val += MOD-b.val; val %= MOD; return *this;}
#define ADDOP(op) constexpr m32 operator op(m32 b){ m32 tmp = *this; return tmp op##= b; }
    ADDOP(+);
    ADDOP(-);
    ADDOP(*);
};
using CC = m32;
constexpr u32 bitceil(u32 x){ return x <= 2 ? 2u : (2u << std::__lg(x-1)); }
namespace Poly{ // }{{{
constexpr CC ROOT = qpow(CC(G), (MOD-1)>>TIM), INVROOT = qpow(ROOT, MOD-2);
template<typename T>
void change(T* arr, int len){
    static std::vector<u32> rev;
    if((int)rev.size() != len){
        rev.resize(len);
        rev[0] = 0;
        UP(i, 1, len){
            rev[i] = rev[i>>1] >> 1;
            if(i&1) rev[i] |= len>>1;
        }
    }
    UP(i, 0, len) if((u32)i < rev[i]) std::swap(arr[i], arr[rev[i]]);
}
void ntt(CC* arr, u32 len, bool idft){
    change(arr, len);
    for(u32 l=2; l<=len; l<<=1){
        CC wn = qpow(idft ? INVROOT : ROOT, (1<<TIM)/l);
        for(u32 st=0; st<len; st+=l){
            CC w = 1;
            UP(j, st, st+l/2){
               CC u = arr[j], v = arr[j+l/2]; v *= w;
               arr[j] = u+v; arr[j+l/2] = u-v;
               w *= wn;
            }
        }
    }
    if(idft){
        CC invlen = qpow(CC(len), MOD-2);
        UP(i, 0u, len) arr[i] *= invlen;
    }
}
using newton_method_func_type = std::function<void(CC*, CC*, u32)>;
using newton_method_edge_type = std::function<CC(CC)>;
void newtons_method_recursive(CC *arr, u32 len, std::vector<CC> &ans, std::vector<CC> &tmp, newton_method_func_type &func, newton_method_edge_type &edge){
    if(len == 1){
        ans.clear();
        ans.push_back(edge(arr[0]));
        return;
    }
    newtons_method_recursive(arr, len/2, ans, tmp, func, edge);
    ans.resize(len*2);
    std::fill(ans.begin()+len/2, ans.end(), 0);
    tmp.resize(len*2);
    std::copy(arr, arr+len, tmp.begin());
    std::fill(tmp.begin()+len, tmp.end(), 0);
    func(ans.data(), tmp.data(), len*2);
}
inline void newtons_method(CC *arr, u32 len, std::vector<CC> &ans, newton_method_func_type &func, newton_method_edge_type &edge){
    std::vector<CC> tmp;
    ans.reserve(len*2);
    tmp.reserve(len*2);
    newtons_method_recursive(arr, len, ans, tmp, func, edge);
    ans.resize(len);
}
void polymul(CC *a, CC *b, u32 len){
    ntt(a, len, false);
    ntt(b, len, false);
    UP(i, 0u, len) a[i] *= b[i];
    ntt(a, len, true);
}
void polyinv(CC *arr, u32 len, std::vector<CC> &ans){
    static newton_method_func_type func = [](CC *a, CC *b, u32 len){ 
        ntt(a, len, false);
        ntt(b, len, false);
        UP(i, 0u, len) a[i] *= (CC(2) - a[i] * b[i]);
        ntt(a, len, true);
    };
    static newton_method_edge_type edge = [](CC a){
        return qpow(a, MOD-2);
    };
    newtons_method(arr, len, ans, func, edge);
}
void integrate(CC *arr, u32 len){
    DOWN(i, len, 1u) arr[i] = arr[i-1] * qpow(CC(i), MOD-2);
    arr[0] = 0;
}
void derivative(CC *arr, u32 len){
    UP(i, 0u, len-1) arr[i] = arr[i+1] * CC(i+1);
    arr[len-1] = 0;
}
void polyln(CC *arr, u32 len, std::vector<CC> &ans){
    std::vector<CC> tmp;
    ans.reserve(len*2);
    tmp.resize(len*2);
    std::copy(arr, arr+len, tmp.begin());
    polyinv(arr, len, ans);
    ans.resize(len*2);
    derivative(tmp.data(), len);
    std::fill(ans.begin()+len, ans.end(), 0);
    std::fill(tmp.begin()+len, tmp.end(), 0);
    polymul(ans.data(), tmp.data(), len*2);
    ans.resize(len);
    integrate(ans.data(), len);
}
void polyexp(CC *arr, u32 len, std::vector<CC> &ans){
    static newton_method_func_type func = [](CC *a, CC *b, u32 len){
        std::vector<CC> ln_a;
        polyln(a, len/2, ln_a);
        ln_a.resize(len);
        std::fill(ln_a.begin()+len/2, ln_a.end(), 0);
        ntt(a, len, false);
        ntt(b, len, false);
        ntt(ln_a.data(), len, false);
        UP(i, 0u, len){
            a[i] *= (CC(1) + b[i] - ln_a[i]);
        }
        ntt(a, len, true);
    };
    static newton_method_edge_type edge = [](CC a){ return 1; };
    newtons_method(arr, len, ans, func, edge);
}
void polyqpow(CC *arr, u32 len, std::vector<CC> &ans, u32 tim, u32 tim_phi, bool k_gt_n){
    // tim % MOD, tim % phi(MOD), tim > len
    u32 min_tim = -1u;
    UP(i, 0u, len) if(!(arr[i] == 0)) {
        min_tim = i;
        break;
    }
    if(min_tim == -1u || min_tim * 1llu * tim >= len || (min_tim && k_gt_n)){
        ans.resize(len);
        std::fill(ans.begin(), ans.end(), 0);
        return;
    }
    u32 ktim = min_tim * tim;
    CC scale = arr[min_tim], invscale = qpow(scale, MOD-2), kscale = qpow(scale, tim_phi);
    u32 nlen = bitceil(len - ktim);
    UP(i, 0u, len-min_tim){
        arr[i] = arr[i+min_tim] * invscale;
    }
    UP(i, len-min_tim, len) arr[i] = 0;
    polyln(arr, nlen, ans);
    UP(i, 0u, nlen) arr[i] = ans[i] * tim;
    polyexp(arr, nlen, ans);
    UP(i, 0u, nlen) ans[i] *= kscale;
    ans.resize(len);
    DOWN(i, len, ktim){
        ans[i] = ans[i-ktim];
    }
    UP(i, 0u, ktim) ans[i] = 0;
}
} // {}}}