#522822 (C++17) No.1145 Sums of Powers

提出ソース

結果

問題	No.1145 Sums of Powers
ユーザー	catupper
提出日時	2020-08-01 01:30:01
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	TLE
実行時間	-
コード長	6,921 bytes
コンパイル時間	2,210 ms
コンパイル使用メモリ	140,924 KB
実行使用メモリ	69,136 KB
最終ジャッジ日時	2023-09-21 05:37:27
合計ジャッジ時間	8,599 ms
ジャッジサーバーID （参考情報）	judge14 / judge15

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	794 ms 31,784 KB
testcase_01	AC	797 ms 31,664 KB
testcase_02	AC	824 ms 32,460 KB
testcase_03	TLE	-
testcase_04	--	-
testcase_05	--	-

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<vector>
#include<map>
#include<set>
#include<string>
#include<queue>
#include<stack>
#include<complex>
#include<numeric>
#include<unordered_map>
using namespace std;
#define MOD 998244353
#define MOD2 998244353
#define INF (1<<29)
#define LINF (1LL<<60)
#define EPS (1e-10)
#define PI 3.1415926535897932384626433832795028
typedef long long Int;
typedef pair<Int, Int> P;
typedef long double Real;
typedef complex<Real> CP;

Int mod_pow(Int x, Int a, Int m = MOD){
    if(a == 0)return 1;
    Int res = mod_pow(x, a / 2, m);
    res = res * res % m;
    if(a % 2)res *= x;
    return res % m;
}

unordered_map<Int, Int> inv_memo;
Int inv(Int x, Int m = MOD){
    Int &res = inv_memo[x];
    if(res != 0)return res;
    return res = mod_pow(x, m-2, m);
}




Int cnt = 0;
void _NTT(vector<Int> &a){
    int n = a.size();
    int other = 0;
    int msb = n/2;
    for(int i = 0;i < n;i++){
        if(other < i)swap(a[i], a[other]);
        for(int b = msb;b > 0;b >>= 1){
            other ^= b;
            if(other & b)break;            
        }        
    }
    for(int i = 1;i*2 <= n;i <<= 1){
        Int w = mod_pow(3, (MOD2-1) / (i*2), MOD2);
        for(int j = 0;j < n;j += i * 2){
            for(Int k = j, ww = 1;k < j + i;k++, ww = ww * w % MOD2){
                auto x = a[k];
                auto y = a[k + i] * ww % MOD;
                a[k] = (x + y) % MOD;
                a[k+i] = (x - y) % MOD;
            }
        }
    }
}

//a.size() must be power of 2
vector<Int> NTT(vector<Int> &a){    
    vector<Int> res = a;
    _NTT(res);
    return res;
}

vector<Int> InvNTT(vector<Int> &a){
    reverse(next(a.begin()), a.end());
    vector<Int> res = NTT(a);
    reverse(next(a.begin()), a.end());
    Int invn = inv(a.size(), MOD2);
    for(int i = 0;i < res.size();i++)res[i] = res[i] * invn % MOD2;
    return res;
}

//a.size() == b.size() == power of 2
vector<Int> conv(vector<Int> &a, vector<Int> &b){
    int n = a.size();
    auto A = NTT(a);
    auto B = NTT(b);
    for(int i = 0;i < n;i++)A[i] = A[i] * B[i] % MOD2;
    auto res = InvNTT(A);
    return res;
}

vector<Int> mult(vector<Int> a, vector<Int> b){
    int n = max(a.size(), b.size());
    int k = 1;
    while(k < n * 2)k *= 2;
    a.resize(k, 0);
    b.resize(k, 0);
    auto res = conv(a, b);
    while(res.back() == 0)res.pop_back();
    return res;
}
vector<Int> add(vector<Int> a, vector<Int> b){
    if(a.size() < b.size())swap(a, b);
    for(int i = 0;i < b.size();i++)(a[i]+=b[i])%=MOD2;
    return a;
}

//a.size() == power of 2, a[0] = 1;
vector<Int> fps_inv(vector<Int> &a){
    Int a0 = a[0];
    Int inva0 = inv(a0);
    int N = a.size();
    for(int i = 0;i < N;i++)(a[i] *= inva0) %= MOD2;
    vector<Int> prea = {a[0], a[1]};
    vector<Int> x = {1, -a[1]};
    int len = 2;
    while(len+1 < N){
        len = len * 2;
        for(int i = prea.size();i < len;i++)prea.push_back(a[i]);
        prea.resize(len*2, 0);
        x.resize(len*2, 0);
        auto tmp = conv(prea, x);
        tmp = conv(tmp, x);
        for(int i = 0;i < len;i++)(x[i] = x[i] * 2 - tmp[i]) %= MOD;
        prea.resize(len);
        x.resize(len);
    }
    a.resize(N);
    for(int i = 0;i < N;i++)(a[i] *= a0) %= MOD2;
    for(int i = 0;i < N;i++)(x[i] *= inva0) %= MOD2;    
    return x;
}

//a.size() == power of 2, a[0] = 1;
vector<Int> fps_log(vector<Int> &a){
    int N = a.size();
    auto da = a;
    for(int i = 0;i < N;i++){
        if(i == N-1)da[i] = 0;
        else da[i] = da[i+1] * (i+1) % MOD2;
    }
    auto x = fps_inv(a);
    da.resize(N*2, 0);
    x.resize(N*2, 0);
    x = conv(x, da);
    for(int i = N-1;i > 0;i--){
        x[i] = x[i-1] * inv(i) % MOD2;
    }
    x[0] = 0;
    x.resize(N);
    return x;
}

//a.size() == power of 2, a[0] = 0;
vector<Int> fps_exp(vector<Int> &a){   
    int N = a.size();
    vector<Int> x = {1, a[1]};
    int len = 2;
    while(len < N){
        len *= 2;
        x.resize(len, 0);
        auto log_x = fps_log(x);
        for(int i = 0;i < len;i++)log_x[i] = a[i] - log_x[i];
        log_x[0] += 1;
        log_x.resize(len * 2, 0);
        x.resize(len * 2, 0);
        x = conv(x, log_x);
        x.resize(len);
    }
    return x;
}

Int mod_rank(Int x, Int y, Int mod){
    int b = 1;
    while(b * b < mod)b++;
    map<Int, Int> giant, baby;
    Int tmp = 1;
    for(int i = 0;i < b;i++){
        baby[tmp] = i;
        tmp = tmp * x % MOD2;
    }
    Int tmp2 = 1;
    for(int i = 0;i < b;i++){
        giant[tmp2] = i;
        tmp2 = tmp * tmp2 % MOD2;
    }
    tmp = 1;
    for(int i = 0;i < b;i++){        
        Int invb = y * inv(tmp, MOD2)  %MOD2;
        if(giant.count(invb)){
            return giant[invb] * b + i;
        }
        tmp = tmp * x % MOD2;
    }
    return -1;
}

Int euclid(Int a, Int b, Int &x, Int &y){
    if(a == 0){
        x = 0;
        y = 1;
        return b;
    }
    Int g = euclid(b % a, a, y, x);
    x -= b/a * y;
    return g;
}

// a / b mod m
Int mod_div(Int a, Int b, Int m){
    Int x, y;
    Int g = euclid(b, m, x, y);
    if(a % g != 0)return -1;
    return x * a / g % m;
}

vector<Int> fps_pow(vector<Int> &a, Int p, Int q = 1){
    Int g = gcd(p, q);p /= g, q /= g;
    Int d = 0;
    Int n = a.size();
    while(d < n && a[d] == 0)d++;
    if(d == n)return a;
    if(d % q != 0)return {};
    vector<Int> tmpa(n, 0);
    for(int i = 0;i+d < n;i++)tmpa[i] = a[i+d];
    Int a0 = tmpa[0];
    for(int i = 0;i < n;i++)(tmpa[i] *= inv(a0))%=MOD2;
    auto l = fps_log(tmpa);
    for(int i = 0;i < n;i++)(l[i] *= p * inv(q) % MOD2)%=MOD2;
    auto ans = fps_exp(l);
    Int r = mod_rank(3, a0, MOD2);if(r == -1)return {};
    Int divr = mod_div(r, q, MOD2-1);if(divr == -1)return {};
    Int a0pow = mod_pow(3, p*divr % (MOD2-1), MOD2);
    for(int i = 0;i < n;i++)(ans[i] *= a0pow) %= MOD2;
    for(int i = n-1;i >= 0;i--){
        Int x = (Int)i - p * d / q;
        if(x < 0)ans[i] = 0;
        else ans[i] = ans[x];
    }
    return ans;
}

// a/b
struct Rational{
    vector<Int> a;
    vector<Int> b;
    Rational(vector<Int> a,vector<Int> b):a(a),b(b){};
    Rational operator+(Rational other){
        auto new_a = add(mult(a, other.b), mult(b, other.a)); 
        auto new_b = mult(b, other.b);
        return Rational(new_a, new_b);
    }
};

int main(){
    int n;
    vector<Rational> a(1 << 18, Rational({0},{1}));
    int m;
    cin >> n >> m;
    for(int i = 0;i < n;i++){
        Int ai;
        cin >> ai;
        a[(1<<17)+i] = Rational({1}, {1, MOD2-ai});
    }
    for(int i = (1<<17)-1;i >= 1;i--){
        a[i] = a[i*2] + a[i*2+1];
    }
    auto up = a[1].a, down = a[1].b;
    int k = 1;
    while(k < m*2)k*=2;
    down.resize(k,0);
    auto ans = mult(up, fps_inv(down));
    for(int i = 1;i <= m;i++){
        if(ans[i] < 0)ans[i] += MOD2;
        cout << ans[i] <<" " ;
    }cout << endl;
    return 0;
}

yukicoder

結果

テストケース

ソースコード