【AtCoder】ARC114C：Sequence Scores（600点）

問題

ARC114C：Sequence Scores

問題の概要

すべての項が\ 0\ である長さ\ N\ の数列\ X\ に対して以下の操作を繰り返し、各項の値が\ 1\ 以上\ M\ 以下の数列\ A\ を作ることを考える。

  • 区間\ [l, r]\ と\ 1\leq m\leq M\ を選択し、各\ l\leq i\leq r\ について、X_i\ を\ \max(X_i, m)\ で置き換える。

\ X\ から\ A\ を作るための最小回数を\ f(A)\ とおいたとき、すべての数列\ A\ に対する\ f(A)\ の総和を\ 998244353\ で割った余りを求めよ。

制約

  • \ 1\leq N, M\leq5000\
  • \ N, M\ は整数

解法

方針

f(A)\ を数列\ A\ を作るためのコストと捉えて問題を考えてみます。

まず、数列の各項に対して１つ１つ左から順番に上記の操作を適用した場合を考えると、明らかに任意の\ A\ に対して、N\ 回の操作で\ A\ を作れることがわかります。作れる数列は\ M^N\ 通りなので、f(A)\ の和は\ NM^N\ 以下です。

最悪でも操作に\ N\ 回しかかからないことがわかりましたが、実際にはもっとコストを削減できる場合もあるはずです。
関数\ f\ を、項数が\ N\ の場合以外にも拡張して考えてみます。
例えば、f(\{3\})\ は明らかに\ 1\ ですが、これに\ 3\ を追加した\ f(\{3, 3\})\ も同様に\ 1\ になります。

つまり、空の数列\ A'\ に\ A\ の項を初項から順番に１つずつ追加していくことを考えたときに、条件次第では追加の前後で\ f(A')\ の値は変わらない、すなわちノーコストで項が追加できることがわかります。

すべての数列\ A\ について、ノーコストで項を追加できたインデックスを数え上げて、NM^N\ から引く方針で答えが出せそうな気がしてきます。

定式化

上の方針で出てきた“ノーコストで追加できる項”の個数を具体的に数式の形で表すことを考えてみましょう。
各\ 1\leq n \leq N, \ 1\leq m \leq M\ に対し、

\ g(n, m)\
　=「\ n\ 番目の項として\ m\ を追加したときに、最小の操作回数が変化しなかったような長さ\ N\ の数列の個数」

と定義します。
”変化しなかった”と過去形で書いているのは、N\ 個の項すべてを追加し終わった後の最終的な数列の個数を数える必要があるためです。

すべての\ n, m\ について\ g(n, m)\ の和をとることにより、削減できるコストを重複・数え漏れなく数え上げることができます。
また、1\ 番目の項を追加する場合は必ずコストがかかるので、n= 1\ の場合は必ず\ g(n, m)= 0\ となり、n\ は\ 2\ 以上のときのみを考えればいいこともわかります。

以上より、関数の和を求める問題を、数列の数え上げの問題へ帰着させることができ、答えを

\begin{aligned} \sum_{A}f(A)=NM^N-\sum_{n=2}^{N}\sum_{m=1}^{M}g(n, m) \end{aligned}

と表すことができました。

ノーコストで追加できる条件

定式化できたところで次に、最小の操作回数が変化しない（＝ノーコストで追加できる）条件を調べてみましょう。
数列をブロックで表すこととして、いくつか例を載せてみます。

例：４番目の要素として、２を追加するとき（\ n=4, m=2\ ）

f:id:tmg_dayo:20220109014227j:plain

上の例を観察すると、\ n\ 個目の要素として\ m\ を追加するときに、次のような条件を満たせばノーコストで追加できることがわかります。

①　ある\ 1\leq i < n\ に対し、\ A_i = m\ が成り立つ
②　①を満たす最大の\ i\ に対して、「\ i < j < n\Longrightarrow A_j > m \ 」が成り立つ

条件を図示すると以下のとおりです。

f:id:tmg_dayo:20220109015345p:plain

上の図の青い部分の数列の選び方の総数を左から順番に計算すると、それぞれ\ M^{i-1}, (M-m)^{n-i-1}, M^{N-n}\ となります。
i\ は\ 1\leq i < n-1\ の範囲を動くことから、g(n, m)\ は次の式で表すことができます。

\begin{align} g(n, m) &= \sum_{i=1}^{n-1}M^{N-n+i-1} \left(M-m\right)^{n-i-1} \\ &= \sum_{i=2}^{n}M^{N-\left(n-i\right)-2} \left(M-m\right)^{n-i} \\ &= \sum_{i=0}^{n-2}M^{N-i-2} \left(M-m\right)^{i} \end{align}

以上の議論より、答えを以下の形で表せることがわかりました。

\begin{align} \sum_{A}f(A)=NM^N-\sum_{n=2}^{N}\sum_{m=1}^{M}\sum_{i=0}^{n-2}M^{N-i-2} \left(M-m\right)^{i} \end{align}

これは\ O(N^2M)\ で計算できます。

高速化

このままではTLEしてしまうので、最後に計算量削減を行います。
先程の式の第２項は以下のように変形できます。

\begin{align} & \sum_{n=2}^{N}\sum_{m=1}^{M}\sum_{i=0}^{n-2}M^{N-i-2} \left(M-m\right)^{i} \\ &= \sum_{n=2}^{N}\sum_{i=0}^{n-2}M^{N-i-2}\left(\sum_{m=1}^{M} \left(M-m\right)^{i}\right) \\ &= \sum_{n=2}^{N}\sum_{i=0}^{n-2}M^{N-i-2}\left(\sum_{m=0}^{M-1} m^{i}\right) \end{align}

ここで、s(i) = \sum_{m=0}^{M-1} m^{i}\ とおくと、これはすべての\ 0 \leq i \leq n-2\ に対して計算しても\ O(NM)\ しかかかりません。

したがって、事前にすべての\ s(i)\ を計算しておくことで、

\begin{align} \sum_{A}f(A)=NM^N-\sum_{n=1}^{N}\sum_{i=0}^{n-2}M^{N-i-2}s(i) \end{align}

と二重和の形で書くことができるため、全体として\ O(NM)\ まで削減できます。これは十分に高速なのでAC可能です。
なお、各式に出てくる累乗についてもその都度求めるのは時間がかかるため、これも事前に計算しておきましょう。

コード（PyPy3）

提出コード（PyPy3：1208ms）
PythonはTLEでした（涙）

def main():
    # 入力、modの定義
    N, M = map(int, input().split())
    mod = 998244353
    # 配列の事前計算（POW[i][j] = pow(i, j) % mod）
    POW = [[1]*5001 for _ in [0]*5001]
    for i in range(5001):
        tmp = 1
        for j in range(1, 5001):
            tmp = tmp*i % mod
            POW[i][j] = tmp
    # s(i)の事前計算
    s = [0]*(N-1)
    for i in range(N-1):
        tmp = 0
        for m in range(M):
            tmp += POW[m][i]
        s[i] = tmp % mod
    # 計算（本体）
    ans = N*POW[M][N] % mod
    for n in range(2, N+1):
        for i in range(0, n-1):
            ans -= POW[M][N-i-2]*s[i]
        ans %= mod
    # 出力
    print(ans)
main()