ABC 137 F - Polynomial Construction(600)

半分寝ながら書いてるのでおかしなところがあったら教えて下さい

問題概要

見づらくなったので省略
リンクから見てください

お気持ち

素直に考えると未知の変数 $p$ 個に式が $p$ 本なので連立方程式解いて $O(p^3)$ 、 $p\leq2999$ なので遅い

$a = \{0, 0, 0, 0, ..., 0, 0\}$ となるような $f(x)$ は容易に構成できるとして(入力例2)、
$a' = \{0, 0, 0, 0, ... , 0, 1, 0, ... , 0\}$ ( $a'_i$ のみ $1$ ) みたいな $g_i(x)$ は作れるかということを考えると、

$g_i(x) := 1 - (x - i)^{p - 1}$
とすると使えそう (フェルマーの小定理より条件を満たすし、展開すると $p-1$ 次の多項式になってうれしい)*1

$a_i = 1$ となるような i について $g_i(x)$ を足し合わせると $a'$ の合計が $a$ と一致するので $f(x)$ が得られる。

$g_i(x)$ を展開して整理すると、

$g_i(x) = 1 - \sum_{j = 0}^{p - 1} (-i)^{p - 1 - j} \binom{p-1}{j} x^j$

これを足し合わせて $f(x)$ を求めたらおわり

二項係数はパスカルの三角形を構築して $O(p^2)$ でも得られてこの場合十分高速だが前計算 $O(p)$ クエリ $O(1)$ でも処理できます([combination 前計算][検索])

コード

実際に提出したコードから不要な部分を取り除いたもの
$O(p^2 log p)$ なんだけど同じ計算量でも落ちている人がいるらしい(?)
このコードもjを逆順に回すことで mod_pow を消すことができて log が落とせる(けど通ったからいいでしょとなってやってない)

#include <bits/stdc++.h>
using namespace std;

#define  range(i, l, r) for (int i = (int)(l); i < (int)(r); ++(i))

template <typename T = int64> T input() { T x; cin >> x; return x; }

template <typename T> ostream &operator<<(ostream &os, vector<T> &v) { range(i, 0, v.size()) { os << v[i] << (i + 1 != v.size() ? " " : ""); } return os; } 
template <typename T> istream &operator>>(istream &is, vector<T> &v) { for (T &in : v) is >> in; return is; }

void solver() {
  int p = input();
  
  vector< int > a(p);
  cin >> a;

// ここからライブラリ
  vector< int > fact(p + 1), finv(p + 1), invs(p + 1);
  invs[1] = 1;
  fact[0] = fact[1] = finv[0] = finv[1] = 1;
  range(i, 2, p + 1) {
    invs[i] = p - invs[p % i] * (p / i) % p;
    fact[i] = 1ll * fact[i - 1] * i % p;
    finv[i] = finv[i - 1] * invs[i] % p;
  }

  auto Ct = [&](int a, int b) {
    if (a < b) return 0;
    return int(1ll * fact[a] * (1ll * finv[b] * finv[a - b] % p) % p);
  };

  auto mod_pow = [&](int a, int x) {
    int64 res = 1;
    for (; x; x >>= 1) {
      if (x & 1) (res *= a) %= p;
      (a *= a) %= p;
    }
    return int(res);
  };
// ここまでライブラリ

  vector< int > b(p);

  range(i, 0, p) {
    if (a[i] == 0) continue;
    
    range(j, 0, p) {
      (b[j] -= 1ll * mod_pow(-i, p - 1 - j) * Ct(p - 1, j) % p) %= p;
      if (b[j] < 0) b[j] += p;
    }

    (b[0] += 1) %= p;
  }

  cout << b << endl;
}

signed main(int argc, char *argv[]) {
  solver();
}