えびちゃんの日記

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

非再帰で補間多項式を求めるよ

問題設定

たかだか \(n\) 次式 \(f(x) = \sum_{i=0}^n a_i x^i\) がありますが、与えられません。

\(n+1\) 個の点 \(x_0, x_1, \dots, x_n\) と、それらの点での \(f\) の値 \(y_0, y_1, \dots, y_n\) が与えられます。 すなわち、\(y_i = f(x_i)\) (\(i=0, 1, \dots, n\)) です。 \(f\) を求めてください。 ただし、各計算は \(998244353 = 119\cdot 2^{23}+1\) を法として行ってください。

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

考察

え〜 Lagrange 補間をしたいです。

\[ l_j(x) = \prod_{\substack{0\le i\le n \\ i\neq j}} \frac{x-x_i}{x_j-x_i}. \] とすると、\(l_j(x_j) = 1\) かつ \(l_j(x_i) = 0\) (\(i\neq j\)) ですから、求めたい多項式は次のよう書けます。 \[ \begin{aligned} f(x) &= \sum_{j=0}^n y_j\cdot l_j(x)\\ &= \sum_{j=0}^n \left(y_j\cdot \prod_{\substack{0\le i\le n \\ i\neq j}} \frac{1}{x_j-x_i}\right) \cdot \prod_{\substack{0\le i\le n \\ i\neq j}} (x-x_i). \end{aligned} \]

係数

まずは、係数の分母である \[ w_j := \prod_{\substack{0\le i\le n \\ i\neq j}} (x_j-x_i) \] の部分を求めることを考えます。 次のように \(l(x)\) を定義します。 \[ l(x) = \prod_{i=0}^n (x-x_i). \] このとき、\(w_j = l'(x_j)\) が成り立ちます。 いわゆる積の微分で項一つずつを微分するのを \(l(x)\) について適用するとわかります。

よって、\(l'(x)\) の \(x_0, x_1, \dots, x_n\) での値を高速に求めればよいです。 最近そういう記事を見ましたね。

rsk0315.hatenablog.com

\(l(x)\) を求めるとき、次のように \(i\) の昇順などで掛けていくと \(\Omega(n^2)\) 時間かかってだめです。

polynomial l = {1};
for (size_t i = 0; i <= n; ++i) l *= polynomial{-x[i], 1};

トーナメント戦のような感じで、隣同士を掛け合わせていくことで \(O(n\log(n)^2)\) 時間を達成できます。 これにより、係数 \(a_j := y_j / w_j\) を \(O(n\log(n)^2)\) 時間で求めることができました。

多項式の構築

あとは、実際に多項式を作るパートです。 やりたいことは、各 \(a_j\) に対して \( (x-x_i)\) (\(i\neq j\)) を掛けて、それらを足し合わせたものを得ることです。

次のように \(f_{l, r}(x)\) を定義します。 \[f_{l, r}(x) = \sum_{j=l}^{r-1} a_j \prod_{\substack{l\le i\lt r \\ i\neq j}} (x-x_i).\] \(f_{0, n}(x)\) が求めたいものです。

まず、基礎ステップとして \(f_{j, j+1}(x) = a_j\) です。 あとは、\(f_{l, c}(x)\) と \(f_{c, r}(x)\) から \(f_{l, r}(x)\) を求めることができればよいです。 これは簡単で、前者に \( (x-x_c)\cdots(x-x_{r-1})\) を掛け、後者に \( (x-x_l)\cdots(x-x_{c-1})\) を掛け、足し合わせればよいです。

掛ける多項式は、前半パートのトーナメントで求めたものを再利用します。

実装

template <typename ModInt>
polynomial<ModInt> interpolate(std::vector<ModInt> const& xs, std::vector<ModInt> const& ys) {
  size_t n = xs.size();
  size_t n2 = ceil2(n);
  std::vector<polynomial<ModInt>> mul(n2+n2, {1}), g(n2+n2);
  for (size_t i = 0; i < n; ++i) mul[n2+i] = {-xs[i], 1};
  for (size_t i = n2; i-- > 1;) mul[i] = mul[i<<1|0] * mul[i<<1|1];

  auto f = mul[1];
  f.differentiate();

  g[1] = f % mul[1];
  for (size_t i = 2; i < n2+n; ++i) g[i] = g[i>>1] % mul[i];
  for (size_t i = 0; i < n; ++i) g[n2+i] = {ys[i] / g[n2+i][0]};
  for (size_t i = n2; i--;) g[i] = g[i<<1|0] * mul[i<<1|1] + g[i<<1|1] * mul[i<<1|0];
  return g[1];
}

みじかいね。 多点評価の実装 同様に非再帰セグ木をイメージしています。

考察を理解して非再帰セグ木状の配列を思い浮かべるとわりと自然だと思うので、詳しい説明はしません。 イメージは多点評価の実装とほぼ同じだと思いますので、説明が必要なら上のリンク先の実装の説明を見てください。

おわり

judge.yosupo.jp

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