F - Minimum Bounding Box 2 Editorial by Kude


\(H\), \(W\) に依らない時間計算量 \(O(K \log K)\) の解法です。

まず簡単に説明します。 選ぶ順番も区別することにすると、 \(K\) マスの選び方は \({}_{HW}P_{K}\) 通りあります。 各選び方に対するスコアの総和(これを \(S\) とします)が求まればよいです。 多項式\(f(x) = x(x - 1)(x - 2) \dots (x - K + 1)\) を用いると \(\displaystyle S = \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} \left( f(HW) - 2 f(aW) - 2 f(Hb) + 4 f(ab) \right)\) と書けます。この計算が \(O(K \log K)\) で行えます。

\(S\) の計算式の導き方

公式解説と同様の包除原理を考えます。座標は0-indexedとします。 \(K\) マス全てを \((a, b)\) より左上側から選ぶような場合の数を \(C_{\text{左上}}(a, b)\) などと書くことにすると、マス \((a, b)\) の寄与は

\[ C_{\text{全体}}(a,b) - C_{\text{上}}(a, b) - C_{\text{下}}(a, b) - C_{\text{左}}(a, b) - C_{\text{右}}(a, b) + C_{\text{左上}}(a, b) + C_{\text{左下}}(a, b) + C_{\text{右上}}(a, b) + C_{\text{右下}}(a, b) \]

と書けるので、

\[ S = \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} \left( C_{\text{全体}}(a,b) - C_{\text{上}}(a, b) - C_{\text{下}}(a, b) - C_{\text{左}}(a, b) - C_{\text{右}}(a, b) + C_{\text{左上}}(a, b) + C_{\text{左下}}(a, b) + C_{\text{右上}}(a, b) + C_{\text{右下}}(a, b) \right) \\ \]

となります。 対称性から \(\displaystyle \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} C_{\text{下}}(a, b) = \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} C_{\text{上}}(a, b)\)\(\displaystyle \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} C_{\text{右下}}(a, b) = \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} C_{\text{左上}}(a, b)\) などが言えるので、

\[ S = \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} \left( C_{\text{全体}}(a,b) - 2 C_{\text{上}}(a, b) - 2 C_{\text{左}}(a, b) + 4 C_{\text{左上}}(a, b) \right) \\ = \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} \left( f(HW) - 2 f(aW) - 2 f(Hb) + 4 f(ab) \right) \]

となります。

\(S\) の計算方法

多項式\(f(x)\)多項式・形式的べき級数-高速に計算できるもの - 第 1 種 Stirling 数 にある通り、分割統治 + Polynomial Taylor Shiftで \(O(K \log K)\) で展開可能です。

\(S\)\(\displaystyle \sum_{a=0}^{H-1} \sum_{b=0}^{W-1} \left( f(HWx) - 2 f(aWx) - 2 f(Hbx) + 4 f(abx) \right)\) の係数の総和です。この多項式も 多項式・形式的べき級数-高速に計算できるもの - \(\sum_{k=L}^{R-1} f(kx)\) の計算 にある通り、\(O(K \log K)\) で計算可能です。 例えば \(\displaystyle \sum_{a=0}^{H-1}\sum_{b=0}^{W-1}f(abx)\) を求めるには、\(\displaystyle g(x) = \sum_{b=0}^{W-1} f(bx)\) を求めたのち、\(\displaystyle \sum_{a=0}^{H-1} g(ax)\) を求めればよいです。あるいは直接

\[ [x^i] \sum_{a=0}^{H-1}\sum_{b=0}^{W-1} f(abx) \\ = \left(\sum_{a=0}^{H-1} \sum_{b=0}^{W-1} a^i b^i \right) [x^i]f(x) \\ = \left(\sum_{a=0}^{H-1} a^i \right) \left(\sum_{b=0}^{W-1} b^i \right)[x^i]f(x) \]

と考えてもよいです。

実装例(C++):

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
using mint = modint998244353;

vector<mint> Fact, FactInv;

void init_fact(int n);
vector<mint> Taylor_shift(vector<mint> f, mint c);        // f(x+c)
vector<mint> Stirling_first_signed(int k);                // x(x-1)...(x-k+1)
vector<mint> fps_inverse(vector<mint> f, int precision);  // 1/f
vector<mint> kth_power_sums(mint n, int m);               // 0^k+1^k+...+(n-1)^k
                                                          // for each k in [0,m)

int main() {
  int h0, w0, k;
  cin >> h0 >> w0 >> k;
  mint h = h0, w = w0;
  init_fact(k + 10);
  vector<mint> xPk = Stirling_first_signed(k),
               sums_h = kth_power_sums(h, k + 1),
               sums_w = kth_power_sums(w, k + 1);
  mint ans;
  mint pow_h = h, pow_w = w;
  for (int i = 0; i <= k; i++) {
    mint v;
    v += pow_h * pow_w;
    v -= 2 * pow_h * sums_w[i];
    v -= 2 * sums_h[i] * pow_w;
    v += 4 * sums_h[i] * sums_w[i];
    ans += xPk[i] * v;
    pow_h *= h;
    pow_w *= w;
  }
  mint hwPk = 1;
  for (int i = 0; i < k; i++) hwPk *= h * w - i;
  ans /= hwPk;
  cout << ans.val() << endl;
}

void init_fact(int n) {
  if (int sz = Fact.size(); sz < n) {
    if (sz == 0) {
      Fact.push_back(1);
      sz++;
    }
    for (; sz < n; sz++) {
      Fact.push_back(Fact.back() * mint::raw(sz));
    }
  }
  if (int sz = FactInv.size(); sz < n) {
    FactInv.resize(n);
    mint x = Fact[n - 1].inv();
    for (int i = n - 1; i >= sz; i--) {
      FactInv[i] = x;
      x *= mint::raw(i);
    }
  }
}

vector<mint> Taylor_shift(vector<mint> f, mint c) {
  int n = f.size();
  init_fact(n);
  for (int i = 0; i < n; i++) {
    f[i] *= Fact[i];
  }
  vector<mint> g(n);
  mint pow_c = 1;
  for (int i = 0; i < n; i++) {
    g[n - 1 - i] = pow_c * FactInv[i];
    pow_c *= c;
  }
  vector<mint> h = convolution(f, g);
  h.erase(h.begin(), h.end() - n);
  for (int i = 0; i < n; i++) {
    h[i] *= FactInv[i];
  }
  return h;
}

vector<mint> Stirling_first_signed(int k) {
  assert(k >= 0);
  if (k == 0) return {1};
  vector<mint> f{0, 1};
  int k_now = 1;
  for (int i = __lg(k) - 1; i >= 0; i--) {
    f = convolution(f, Taylor_shift(f, -k_now));
    k_now *= 2;
    if (k >> i & 1) {
      f = convolution(f, {-k_now, 1});
      k_now++;
    }
  }
  assert(k_now == k);
  return f;
}

vector<mint> fps_inverse(vector<mint> f, int precision) {
  int n = f.size();
  assert(n >= 1 && f[0] != 0);
  int len = 1;
  const int z = 1 << internal::ceil_pow2(precision);
  vector<mint> g{f[0].inv()};
  const mint inv4 = mint(4).inv();
  mint inv4k = 1;
  while (len < z) {
    int nlen = 2 * len;
    vector<mint> ft(f.begin(), f.begin() + min(n, nlen));
    ft.resize(nlen);
    butterfly(ft);
    vector<mint> gt = g;
    gt.resize(nlen);
    internal::butterfly(gt);
    for (int i = 0; i < nlen; i++) ft[i] *= gt[i];
    internal::butterfly_inv(ft);
    ft.erase(ft.begin(), ft.begin() + len);
    ft.resize(nlen);
    internal::butterfly(ft);
    for (int i = 0; i < nlen; i++) ft[i] *= gt[i];
    internal::butterfly_inv(ft);
    inv4k *= inv4;
    mint c = -inv4k;
    for (int i = 0; i < len; i++) g.emplace_back(c * ft[i]);
    len = nlen;
  }
  g.resize(precision);
  return g;
}

vector<mint> kth_power_sums(mint n, int m) {
  // e^0+e^x+e^2x+...+e^(n-1)x
  // = ((e^nx-1)/x) / ((e^x-1)/x)
  // =: f / g
  init_fact(m + 1);
  vector<mint> f(m), g(m, 1);
  mint pow_n = n;
  for (int i = 0; i < m; i++) {
    f[i] = pow_n * FactInv[i + 1];
    g[i] = FactInv[i + 1];
    pow_n *= n;
  }
  vector<mint> res = convolution(f, fps_inverse(g, m));
  res.resize(m);
  for (int i = 0; i < m; i++) res[i] *= Fact[i];
  return res;
}

提出 (850 ms): https://atcoder.jp/contests/abc297/submissions/40582850

posted:
last update: