結果

問題 No.1068 #いろいろな色 / Red and Blue and more various colors (Hard)
コンテスト
ユーザー tonyu0
提出日時 2021-02-07 21:21:44
言語 C++17
(gcc 13.3.0 + boost 1.89.0)
結果
AC  
実行時間 1,066 ms / 3,500 ms
コード長 2,799 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 1,644 ms
コンパイル使用メモリ 121,228 KB
最終ジャッジ日時 2025-01-18 14:24:24
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 29
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

#include <algorithm>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;
using ll = long long;
#define rep(i, j, n) for (ll i = j; i < (n); ++i)
#define rrep(i, j, n) for (ll i = (n)-1; j <= i; --i)

template <typename T>
std::ostream& operator<<(std::ostream& os, std::vector<T>& a) {
  for (size_t i = 0; i < a.size(); ++i) os << (i > 0 ? " " : "") << a[i];
  return os << '\n';
}

[[maybe_unused]] constexpr ll MOD = 998244353;
[[maybe_unused]] constexpr int INF = 0x3f3f3f3f;
[[maybe_unused]] constexpr ll INFL = 0x3f3f3f3f3f3f3f3fLL;

// TODO: 多項式同士の演算として書く
template <long long mod = 998244353>
class modFFT {
public:
  modFFT() {}

  vector<long long> multiply(vector<long long> a, vector<long long> b) {
    int an = a.size();
    int bn = b.size();
    int need = an + bn - 1;
    int nn = 1;
    while (nn < 2 * an || nn < 2 * bn) nn <<= 1;
    a.resize(nn);
    b.resize(nn);
    fft(a);
    fft(b);
    for (int i = 0; i < nn; i++) a[i] = a[i] * b[i] % mod;
    reverse(++a.begin(), a.end());
    fft(a);
    long long inv = pow(nn, mod - 2);
    for (int i = 0; i < nn; i++) a[i] = a[i] * inv % mod;
    a.resize(need);
    return a;
  }
  vector<long long> pow(vector<long long> a) {
    size_t n = a.size();
    if (n == 1) return {1, a[0] - 1};

    vector<long long> b(a.begin(), a.begin() + n / 2);
    vector<long long> c(a.begin() + n / 2, a.end());
    return multiply(pow(b), pow(c));
  }

private:
  long long pow(long long a, long long b) {
    long long x = 1, step = 1 << 30;
    while (step > 0) {
      x = x * x % mod;
      if (step & b) x = x * a % mod;
      step >>= 1;
    }
    return x;
  }

  void fft(vector<long long>& a) {
    int n = a.size();
    for (int i = 0; i < n; i++) {
      int j = 0;
      int x = i, y = n - 1;
      while (y > 0) {
        j = (j << 1) + (x & 1);
        x >>= 1;
        y >>= 1;
      }
      if (i < j) swap(a[i], a[j]);
    }
    for (int len = 1; len < n; len *= 2) {
      long long root = pow(5, (mod - 1) / (2 * len));
      for (int i = 0; i < n; i += 2 * len) {
        long long w = 1;
        for (int j = 0; j < len; j++) {
          long long u = a[i + j];
          long long v = a[i + j + len] * w % mod;
          a[i + j] = (u + v) % mod;
          a[i + j + len] = (mod + u - v) % mod;
          w = w * root % mod;
        }
      }
    }
  }
};

int main() {
  cin.tie(0)->sync_with_stdio(0);
  int n, q, b;
  cin >> n >> q;
  vector<ll> a(n);
  rep(i, 0, n) cin >> a[i];
  modFFT fft = modFFT();
  vector<ll> c = fft.pow(a);

  reverse(c.begin(), c.end());
  rep(i, 0, q) {
    cin >> b;
    cout << c[b] << '\n';
  }
  return 0;
}
0