test/fps/multpoint_evalution_geo.test.cpp

View this file on GitHub
Last update: 2024-09-05 20:11:34+09:00
Problem: https://judge.yosupo.jp/problem/multipoint_evaluation_on_geometric_sequence

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/multipoint_evaluation_on_geometric_sequence"
#include "../../fps/Multipoint_Evaluation_Geo.hpp"

#include <iostream>
int main(){
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    using mint = atcoder::modint998244353;

    int N, M, a, r;
    std::cin >> N >> M >> a >> r;
    std::vector<mint> c(N);
    int b;
    for (auto &x : c){
        std::cin >> b;
        x = b;
    }
    auto ans = po167::Multipoint_Evaluation_Geo(c, (mint)(a), (mint)(r), M);
    for (int i = 0; i < (int)ans.size(); i++){
        std::cout << ans[i].val() << (i + 1 == (int)ans.size() ? "\n" : " ");
    }
}

#line 1 "test/fps/multpoint_evalution_geo.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/multipoint_evaluation_on_geometric_sequence"
#line 2 "fps/FPS_cyclic_convolution.hpp"
#include <vector>
#include <atcoder/convolution>

namespace po167{
// |f| = |g| = 2 ^ n
template<class T>
std::vector<T> FPS_cyclic_convolution(std::vector<T> f, std::vector<T> g){
    atcoder::internal::butterfly(f);
    atcoder::internal::butterfly(g);
    for (int i = 0; i < (int)f.size(); i++) f[i] *= g[i];
    atcoder::internal::butterfly_inv(f);
    T iz = (T)(1) / (T)(f.size());
    for (int i = 0; i < (int)f.size(); i++) f[i] *= iz;
    return f;
}
}
#line 2 "fps/Multipoint_Evaluation_Geo.hpp"
namespace po167{
// return {f(ar^0), f(ar^1), f(ar^2), ... f(ar^{len - 1})}
template <class T>
std::vector<T> Multipoint_Evaluation_Geo(
    std::vector<T> f,
    T a, T r, int len = -1
){
    if (len == -1) len = f.size();
    if (r == 0){
        T tmp = 1;
        std::vector<T> res(len, f[0]);
        for (int i = 1; i < (int)f.size(); i++){
            tmp *= a;
            res[0] += tmp * f[i];
        }
        return res;
    }
    /*
    X[k] = sum (fi * (ar^k)^i)
         = sum (fi * a^i * r^{ki})
         = sum (fi * a^i * r^t[k + i] * r^-t[k] * r^-t[i])
         = r^-t[k] sum r^t[k + i](fi * a^i * r^-t[i])
    */
    int n = f.size();
    T invr = (T)(1) / (T)(r);
    T tmp = 1, tmpr1 = 1, tmpr2 = 1;
    for (int i = 0; i < n; i++){
        f[i] *= tmp * tmpr2;
        tmpr2 *= tmpr1;
        tmpr1 *= invr;
        tmp *= a;
    }
    std::vector<T> g(n + len - 1);
    tmpr1 = 1, tmpr2 = 1;
    for (int i = 0; i < n + len - 1; i++){
        g[i] = tmpr2;
        tmpr2 *= tmpr1;
        tmpr1 *= r;
    }
    std::reverse(f.begin(), f.end());
    int size = 1;
    while (size < n + len - 1) size *= 2;
    f.resize(size);
    g.resize(size);
    f = po167::FPS_cyclic_convolution(f, g);
    std::vector<T> res(len);
    for (int i = n - 1; i < n + len - 1; i++){
        res[i - n + 1] = f[i];
    }
    tmpr1 = 1, tmpr2 = 1;
    for (int i = 0; i < len; i++){
        res[i] *= tmpr2;
        tmpr2 *= tmpr1;
        tmpr1 *= invr;
    }
    return res;
}
}
#line 3 "test/fps/multpoint_evalution_geo.test.cpp"

#include <iostream>
int main(){
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    using mint = atcoder::modint998244353;

    int N, M, a, r;
    std::cin >> N >> M >> a >> r;
    std::vector<mint> c(N);
    int b;
    for (auto &x : c){
        std::cin >> b;
        x = b;
    }
    auto ans = po167::Multipoint_Evaluation_Geo(c, (mint)(a), (mint)(r), M);
    for (int i = 0; i < (int)ans.size(); i++){
        std::cout << ans[i].val() << (i + 1 == (int)ans.size() ? "\n" : " ");
    }
}