ARC114-C Sequence Scores 解説 - ふろんてぃあ

問題リンク

C - Sequence Scores

問題概要

$1$ 以上 $M$ 以下の整数からなる長さ $N$ の数列 $A$ に対して、 $f(A)$ を、すべての要素が $0$ である長さ $N$ の数列 $X$ に以下の操作を繰り返して $A$ を得るのに必要な最小の操作回数とする。

整数 $l, r\ (1\le l \le r \le N)$ と、正整数 $v \le M$ を1つ選び、各 $X_i\ (l \le i \le r)$ を $\max(X_i, v)$ で置き換える。

ありうる $M^N$ 通りの $A$ すべてに関する $f(A)$ の和を $\mod 998244353$ で求めよ。

思考過程

こういう「ありうるすべての○○についての総和」みたいな問題は、まとめて数え上げられることの方が多い気がする。 $A_i$ の位置に $x$ を追加したとき、操作回数が増えるのは、

これまでに $x$ が使われていない場合。
これまでに $x$ が使われたが、その後 $x$ より小さい数字が使われている場合

の2パターンに分けられる。 $i, x$ を固定し、それぞれを丁寧に数える。

パターン1

$A_i = x$ かつ、すべての $j \lt i$ について $A_j \neq x$ であるような $A$ の数は、

$(M-1)^{i-1}M^{N-i}$

に等しい。

パターン2

ある $k \lt i$ が存在して $A_k = x$ であり、かつ、ある $l\ (k \lt l \lt i)$ が存在して $A_l \lt x$ となるような $A$ の数を求める。 $A_j = x\ (j \lt i)$ である最大の $j$ の値（これを $k$ とおく）で場合分けをする。 $k$ を固定したとき、

$j \lt k, i \lt j$ については $A_j$ は自由に決められる。
$A_k,\ A_i$ は仮定より $x$ に等しい。
$k \lt j \lt i$ については、「『すべての $j$ について $A_j \ge x$ 』でなく、かつすべての $j$ について $A_j \neq x$ である」ものを数えればいい。

また、これらは互いに独立に決められるので、 $k$ を固定したときのパターン数は、

$\displaystyle{M^{k-1}\cdot ((M-1)^{i - k - 1} - (M - x) ^ {i - k - 1})\cdot M^{N - i}}$

に等しい。

解法

あとは上の2式の $i, x, k$ に関する総和を取ればよい。このままだと $O(N^2 M)$ かかってしまうが、 $k$ に関する和は等比数列の和の公式を使って以下のようにまとめられるので、 $O(NM)$ で解ける。

$\begin{aligned} &\sum_{k = 1}^{i - 2} M^{k-1}\cdot ((M-1)^{i - k - 1} - (M - x) ^ {i - k - 1})\\ &= \sum_{k = 1}^{i - 2} \frac{M^k}{M} \cdot \left( \frac{(M-1)^{i-1}}{(M-1)^{k}} -\frac{(M-x)^{i-1}}{(M-x)^k} \right) \\ &= \sum_{k = 1}^{i - 2} \left(\frac{(M-1)^{i-1}}{M} \cdot \left(\frac{M}{M-1}\right)^k - \frac{(M-x)^{i-1}}{M} \cdot\left(\frac{M}{M-x}\right)^k \right) \\ &= \frac{(M-1)^{i-1}}{M} \cdot \frac{M}{M-1} \cdot \frac{1 - \left(\frac{M}{M-1}\right)^{i-2}}{1 - \frac{M}{M-1}} - \frac{(M-x)^{i-1}}{M} \cdot\frac{M}{M-x}\cdot\frac{1 - \left(\frac{M}{M-x}\right)^{i-2}}{1 - \frac{M}{M-x}} \\ &= (M - 1)^{i - 2}\cdot \frac{1 - \left(\frac{M}{M-1}\right)^{i-2}}{1 - \frac{M}{M-1}} - (M-x)^{i-2} \cdot \frac{1 - \left(\frac{M}{M-x}\right)^{i-2}}{1 - \frac{M}{M-x}} \end{aligned}$

実装

TLがやや厳しめで、毎回累乗や逆元を計算していると間に合わないので、前計算が必要だった（1146ms）。

#include <bits/stdc++.h>
#include <atcoder/modint>

using namespace std;
using namespace atcoder;
using mint = modint998244353;

int main(){
    int n, m; cin >> n >> m;

    if(m == 1){
        cout << 1 << '\n';
        return 0;
    }

    //逆元の前計算
    vector<mint> inv(5010, 0);
    for(int i = 1; i < 5010; i++){
        inv[i] = mint(i).inv();
    }

    //累乗の前計算
    vector<vector<mint>> pow(5010, vector<mint>(5010, 0));
    for(int i = 0; i < 5010; i++){
        mint p = 1;
        for(int j = 0; j < 5010; j++){
            pow[i][j] = p;
            p *= i;
        }
    }

    //1 - (m/(m-i))^jの前計算
    vector<vector<mint>> r(5010, vector<mint>(5010));
    for(int i = 1; i < m; i++){
        mint t = 1, u = inv[m-i] * m;
        for(int j = 0; j <= n; j++){
            r[i][j] = 1 - t;
            t *= u;
        }
    }
    for(int j = 0; j <= n; j++) r[m][j] = 1;

    //r[i][j]の逆元の前計算
    vector<mint> rInv(5010);
    for(int i = 0; i < 5010; i++){
        if(r[i][1] == 0) continue;
        rInv[i] =  r[i][1].inv();
    }
    
    mint ans = 0;
    for(int i = 1; i <= n; i++){
        for(int x = 1; x <= m; x++){
            ans += pow[m-1][i-1] * pow[m][n-i];
            if(i <= 2) continue;
            ans += pow[m-1][i-2] * r[1][i-2] * rInv[1] * pow[m][n-i];
            ans -= pow[m-x][i-2] * r[x][i-2] * rInv[x] * pow[m][n-i];
        }
    }

    cout << ans.val() << '\n';
}

感想

考察自体はすんなり行けたが、添字合わせや定数倍高速化でものすごく時間を取られた。