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: