えびちゃんの日記

えびちゃん(競プロ)の日記です。

非再帰で多項式の多点評価をするよ

問題設定

多項式 \(f(x) = \sum_{i=0}^n a_i x^i\) と、\(m\) 個の異なる点 \(x_0, x_1, \dots, x_{m -1}\) が与えられます。 これらの点での値 \(f(x_0), f(x_1), \dots, f(x_{m -1})\) を求めてください。 ただし、各計算は \(998244353 = 119\cdot 2^{23}+1\) を法として行ってください。

  • \(1\le n\le 2^{17} = 131072\)
  • \(1\le m\le 2^{17} = 131072\)

観察

\(f(x_j)\) は \(f(x)\bmod{(x-x_j)}\) と等しいです。 \[ f(x) = Q(x)\cdot (x-x_j) + f(x_j). \] \(f(x)\bmod{(x-x_j)}\) を高速に求められるとうれしいのですが、これは \(\Omega(n)\) 時間かかってしまい、つらいです。

一般には、\(f(x)\cdot{g(x)}\) を \(M(n)\) 時間で求められるとき、\(O(M(n))\) 時間で \(f(x)\bmod{g(x)}\) を求められることが知られています*1。 これより、\(n\) 次式を \(m\) 次式で割るのを \(O( (n+m)\log(n+m))\) 時間で行えます*2

解法

分割統治で行うことを考えます。 はじめ \([l, r) = [0, m)\) とし、\(j\in[l, r)\) に関して \(f(x_j)\) を求めます。

\(f(x)\) を \( (x-x_l)\cdots(x-x_{r-1})\) で割ります。 \[ f(x) = Q_{l, r}(x)\cdot (x-x_l)\cdots (x-x_{r-1}) + R_{l, r}(x). \] \(f(x_j) = R_{l, r}(x_j)\) であることに注意します。\(R_{l, r}(x)\) はたかだか \(r-l\) 次式です。

\(c = \lfloor \frac{l+r}{2}\rfloor\) とし、\(R_{l, r}(x)\) を \( (x-x_l)\cdots(x-x_{c-1})\) と \( (x-x_c)\cdots(x-x_{r-1})\) でそれぞれ割ります。 \[ \begin{aligned} R_{l, r}(x) &= Q_{l, c}(x)\cdot (x-x_l)\cdots (x-x_{c-1}) + R_{l, c}(x),\\ R_{l, r}(x) &= Q_{c, r}(x)\cdot (x-x_c)\cdots (x-x_{r-1}) + R_{c, r}(x). \end{aligned} \] このとき、\(j\in[l, c)\) については \(f(x_j) = R_{l, r}(x_j) = R_{l, c}(x_j)\) で、\(j\in[c, r)\) については \(f(x_j) = R_{l, r}(x_j) = R_{c, r}(x_j)\) です。 よって、\(r-l\) 次式についての問題一つを \(\frac{r-l}{2}\) 次式の問題二つに帰着できました。

\(R_{l, r}(x)\) から \(R_{l, c}(x)\)、すなわち \(R_{l, r}(x)\bmod{(x-x_l)\cdots (x-x_{c-1})}\) を求めるのは \(O( (r-l)\log(r-l) )\) 時間でできます。 \(R_{c, r}(x)\) についても同様ですから、全体で \(O(m\log(m)^2)\) 時間で求めることができます。

実装

template <typename ModInt>
class polynomial {
public:
  using size_type = size_t;
  using value_type = ModInt;

  // 乗算とか除算とかの実装は省略

  // xs の各要素に対しての f(xs) を返すよ
  std::vector<value_type> multipoint_evaluate(std::vector<value_type> const& xs) const {
    size_type m = xs.size();
    size_type m2 = ceil2(m);  // m 以上の最小の 2 べき
    std::vector<polynomial> g(m2+m2, {1});
    for (size_type i = 0; i < m; ++i) g[m2+i] = {-xs[i], 1};  // (x - xi)
    for (size_type i = m2; i-- > 1;) g[i] = g[i<<1|0] * g[i<<1|1];

    g[1] = (*this) % g[1];
    for (size_type i = 2; i < m2+m; ++i) g[i] = g[i>>1] % g[i];
    std::vector<value_type> ys(m);
    for (size_type i = 0; i < m; ++i) ys[i] = g[m2+i][0];
    return ys;
  }
};

みじかいね。

再帰セグ木のイメージで配列を使います。 今回は 2 べきに丸めた方が実装が楽になるはずなので、そうしています。

g[m2+j] には最初 \( (x-x_j)\) を入れて*3、掛け算でマージしていきます(前半の降順ループ)。 ここまでで g[1] は \( (x-x_0)\cdots(x-x_{m -1})\) です。

次に、g[1] を \(R_{0, m}(x)\) にし、昇順ループをしながら \(R_{l, r}(x)\) を求めていきます。 \(f(x_j)=R_{j, j+1}(x)\) は g[m2+j] にあるので、それを取り出して終了です。

おわり

judge.yosupo.jp

みんなも実装してみてね。

*1:たぶん Newton 法をがんばるといいです。

*2:ナイーブな実装の方が高速そうな \(n, m\) のときは適宜そうしてください。あくまで \(O\) は漸近的なやつなので。

*3:padding が必要なら \(1\) を入れます。