AtCoder Regular Contest 111 F - Do you like query problems?

めちゃくちゃ計算した．~~こんなんどうやったら時間内に解けるんだろう？~~ クエリの読み替えが本質だったっぽくて，これに気付くことが出来ればあるいは...？みたいな感じかもしれない．ただ，今の自分では厳しい...

問題

解法

以降，操作対象の数列を $A$ とし，0-indexed とする．

列全体の動作を考えるのは難しいので，数列の添え字 $t=0,\ldots,N-1$ を固定して要素 1 つ 1 つに注目する．クエリの範囲は全部で $\dfrac{N(N+1)}{2}$ 通りあるが，このうち $t$ が範囲内に含まれるのは $l=0,\ldots,t$ かつ $r=t,\ldots,N-1$ のときで，即ち $(t+1)(N-t)$ 通り．これらの数を毎回書くと大変なので， $I_t=(t+1)(N-t)$ , $C=\dfrac{N(N+1)}{2}$ とおいておく．

解法 (泥臭いので読み飛ばしてもらってもよいです)

現在の $A_t$ の値を持てば DP が出来そうなので，DP テーブルを次のように定義してみる．

$\displaystyle \begin{aligned} X_t[i][x]=&\text{ クエリを $i$ 個処理した後に $A_t=x$ であるような全ての場合の，}\\ &\text{$A_t$ の寄与の総和} \end{aligned}$

$X_t[0]$ に関して，全ての $x$ に対して $X_t[0][x]=0$ である．

$i+1$ 個目のクエリが $\max$ ， $\min$ ， $\mathrm{sum}$ である場合のそれぞれについて， $X_t[i+1][x]$ への遷移を考える (集める DP)．

$\max$ クエリ

$t$ がクエリの範囲内であるかどうかで更に場合分けをする．

$t$ が範囲内の場合

$t$ が範囲内となるような範囲の選び方は $I_t$ 通りである．

$\max$ クエリにおいて $A_t=x$ となるのは，

クエリの前は $A_t=0,\ldots,x-1$ で， $v=x$ のクエリが飛んでくる場合
クエリの前は $A_t=x$ で， $v=0,\ldots,x$ のクエリが飛んでくる場合

の 2 つのケースが考えられる．従って，このケースの寄与は次の式 $(1)$ で表される．

$\displaystyle I_t\left(\sum_{y=0}^{x-1}X_t[i][y]+(x+1)\cdot X_t[i][x]\right)\tag{1}$

$t$ が範囲外の場合

$t$ が範囲内となるような範囲の選び方は $C-I_t$ 通りである．

$t$ が範囲外の場合は元から $A_t=x$ である必要があり， $v$ の値は関係ない．従って，このケースの寄与は次の式 $(2)$ で表される．

$\displaystyle (C-I_t)\cdot M\cdot X_t[i][x]\tag{2}$

$\min$ クエリ

$t$ がクエリの範囲内であるかどうかで更に場合分けをする．

$t$ が範囲内の場合

$t$ が範囲内となるような範囲の選び方は $I_t$ 通りである．

$\min$ クエリにおいて $A_t=x$ となるのは，

クエリの前は $A_t=x,\ldots,M-1$ で， $v=x$ のクエリが飛んでくる場合
クエリの前は $A_t=x$ で， $v=x+1,\ldots,M-1$ のクエリが飛んでくる場合

の 2 つのケースが考えられる．従って，このケースの寄与は次の式 $(3)$ で表される．

$\displaystyle I_t\left(\sum_{y=x}^{M-1}X_t[i][y]+(M-x-1)\cdot X_t[i][x]\right)\tag{3}$

$t$ が範囲外の場合

$\max$ クエリにおいて $t$ が範囲外の場合と全く同様であり，このケースの寄与は次の式 $(4)$ で表される．

$\displaystyle (C-I_t)\cdot M\cdot X_t[i][x]\tag{4}$

$\mathrm{sum}$ クエリ

まず， $\mathrm{sum}$ クエリにおいて $A_t$ の値は変化しないので，任意の範囲に対して $X_t[i][x]$ からの寄与があり， $X_t[i][y]$ ( $y\neq x$ ) からの寄与はない．従って，この場合の寄与は式 $(5)$ で表される．

$\displaystyle C\cdot X_t[i][x]\tag{5}$

また， $\mathrm{sum}$ クエリにおいてはクエリ自体による寄与が存在する． $t$ が範囲内である必要があるので，これは以下のように書ける．

$\displaystyle x\cdot I_t\cdot \text{(クエリを $i$ 個処理した後に $A_t=x$ であるような場合の数)}$

ここで， $\text{(クエリを $i$ 個処理した後に $A_t=x$ であるような場合の数)}$ が分からないので別で求める必要がある．

取り敢えずこの値を $Y_t[i][x]$ とおけば，クエリ自体の寄与は次の式 $(6)$ で表される．

$\displaystyle x\cdot I_t\cdot Y_t[i][x]\tag{6}$

以上 $(1),\ldots,(6)$ より， $X_t[i+1][x]$ の表示は式 $(7)$ のようになる．

$\displaystyle \require{color} \begin{align} X_t[i+1][x]&= \underset{\textstyle\text{$\max$ クエリの寄与}}{\textcolor{red}{\underline{\textcolor{black} { I_t\left(\sum_{y=0}^{x-1}X_t[i][y]+(x+1)\cdot X_t[i][x]\right)+(C-I_t)\cdot M\cdot X_t[i][x] }}}}\\ &+ \underset{\textstyle\text{$\min$ クエリの寄与}}{\textcolor{blue}{\underline{\textcolor{black} { I_t\left(\sum_{y=x}^{M-1}X_t[i][y]+(M-x-1)\cdot X_t[i][x]\right)+(C-I_t)\cdot M\cdot X_t[i][x] }}}}\\ &+ \underset{\textstyle\text{$\mathrm{sum}$ クエリの寄与}}{\textcolor{green}{\underline{\textcolor{black} { C\cdot X_t[i][x]+x\cdot I_t\cdot Y_t[i][x] }}}}\\ &=I_t\sum_{y=0}^{M-1}X_t[i][y]+( (2M+1) \cdot C-M \cdot I_t)\cdot X_t[i][x]+x\cdot I_t\cdot Y_t[i][x] \tag{7} \end{align}$

そして，この問題で求めたいのは $X_t[i][x]$ を $x$ について和をとったものなので， $S_t$ を次のように定義する．

$\displaystyle S_t[i]=\sum_{x=0}^{M-1}X_t[i][x]\tag{8}$

式 $(7)$ の両辺を $x$ について和をとって，式 $(9)$ を得る．

$\displaystyle S_t[i+1]=(2M+1)\cdot C\cdot S_t[i]+I_t\sum_{x=0}^{M-1}x\cdot Y_t[i][x]\tag{9}$

良い感じになってきたが， $Y_t$ が分からないのでこれ以上は計算が進まない．そこで，以下は $Y_t$ を求める．

まず， $Y_t[0]$ は定義より次のようになる．

$\displaystyle Y_t[0][x]=\begin{cases} 1&\text{(if $x=0$)}\\ 0&\text{(otherwise)} \end{cases}$

そして， $Y_t$ も $X_t$ とほぼ同様の漸化式を立てることが出来る．導出過程は省略し，結果を式 $(10)$ に示す．式 $(7)$ との違いは最終項の有無のみである．

$\displaystyle \begin{align} Y_t[i+1][x]&= \underset{\textstyle\text{$\max$ クエリの寄与}}{\textcolor{red}{\underline{\textcolor{black} { I_t\left(\sum_{y=0}^{x-1}Y_t[i][y]+(x+1)\cdot Y_t[i][x]\right)+(C-I_t)\cdot M\cdot Y_t[i][x] }}}}\\ &+ \underset{\textstyle\text{$\min$ クエリの寄与}}{\textcolor{blue}{\underline{\textcolor{black} { I_t\left(\sum_{y=x}^{M-1}Y_t[i][y]+(M-x-1)\cdot Y_t[i][x]\right)+(C-I_t)\cdot M\cdot Y_t[i][x] }}}}\\ &+ \underset{\textstyle\text{$\mathrm{sum}$ クエリの寄与}}{\textcolor{green}{\underline{\textcolor{black} { C\cdot Y_t[i][x] }}}}\\ &=I_t\sum_{y=0}^{M-1}Y_t[i][y]+( (2M+1) \cdot C-M \cdot I_t)\cdot Y_t[i][x] \tag{10} \end{align}$

ここで， $Y_t[i]$ の総和は実は問題文にその答えが書かれており，

$\displaystyle \sum_{x=0}^{M-1}Y_t[i][x]=( (2M+1)\cdot C)^i \tag{11}$

である (もちろん普通に求めてもよい)．従って，更に式 $(10)$ の計算を進めることができ，

$\displaystyle Y_t[i+1][x]=I_t( (2M+1)\cdot C)^i +( (2M+1) \cdot C-M \cdot I_t)\cdot Y_t[i][x]$

となる．これは以下のようにして $x$ について独立に解くことができる．

$\displaystyle \frac{Y_t[i+1][x]}{(a-b_t)^{i+1}}=\frac{Y_t[i][x]}{(a-b_t)^{i}}+\frac{I_t}{a-b_t}\left(\frac{a}{a-b_t}\right)^i\quad \left(\begin{aligned} a &= (2M+1)\cdot C,\\ b_t &= M \cdot I_t. \end{aligned}\right)$

$\displaystyle \begin{aligned} Y_t[i][x]&= (a-b_t)^i\left(Y_t[0][x]+\frac{I_t}{a-b_t}\sum_{j=0}^{i-1}\left(\frac{a}{a-b_t}\right)^j\right)\\ &=(a-b_t)^i\cdot Y_t[0][x]+\frac{1}{M}\left(a^i-(a-b_t)^i\right)\\ &=\begin{cases} \dfrac{1}{M}\left(a^i-(a-b_t)^i\right)+(a-b_t)^i & \text{(if $x=0$)}\\ \dfrac{1}{M}\left(a^i-(a-b_t)^i\right)& \text{(otherwise)} \end{cases} \end{aligned}$

さて， $Y_t$ の具体的な表示が得られたので式 $(9)$ に代入する．

$\displaystyle \begin{align} S_t[i+1] &=a\cdot S_t[i]+I_t\sum_{x=0}^{M-1}x\cdot Y_t[i][x]\\ &=a\cdot S_t[i]+I_t\sum_{x=\textcolor{red}{1}}^{M-1}x\cdot Y_t[i][x]\\ &=a\cdot S_t[i]+I_t\sum_{x=1}^{M-1}x\cdot \dfrac{1}{M}\left(a^i-(a-b_t)^i\right)\\ &=a\cdot S_t[i]+\dfrac{I_t}{M}\left(a^i-(a-b_t)^i\right)\sum_{x=1}^{M-1}x\\ &=a\cdot S_t[i]+\dfrac{I_t\cdot (M-1)}{2}\left(a^i-(a-b_t)^i\right) \end{align}$

この形の漸化式は両辺を $a ^ {i+1}$ で割ると解けて，

$\displaystyle \begin{align} \frac{S_t[i]}{a^{i}} &=\frac{S_t[i-1]}{a^{i-1}}+\frac{I_t\cdot (M-1)}{2a}\left(1-\left(\frac{a-b_t}{a}\right)^{i-1}\right)\\ &=S_t[0]+\frac{I_t\cdot (M-1)}{2a}\sum_{j=0}^{i-1}\left(1-\left(\frac{a-b_t}{a}\right)^j\right)\\ &=\frac{I_t\cdot (M-1)}{2a}\left(i+\frac{a}{b_t}\left(\left(\frac{a-b_t}{a}\right)^i-1\right)\right)\quad(\because X_t[0][x]=0,\;\forall x),\\ S_t[i]&=\frac{I_t\cdot (M-1)}{2}\left(i\cdot a^{i-1}+\frac{(a-b_t)^i-a^i}{b_t}\right)\tag{12} \end{align}$

となる．式 $(12)$ より，求める答えは $(13)$ である．

$\displaystyle \sum_{t=0}^{N-1}S_t[Q]=\sum_{t=0}^{N-1}\frac{I_t\cdot (M-1)}{2}\left(Q\cdot a^{Q-1}+\frac{(a-b_t)^Q-a^Q}{b_t}\right)\tag{13}$

ここで，

$\displaystyle \begin{gather} C=\frac{(N+1)N}{2},& I_t=(t+1)(N-t),\\ a=(2M+1)\cdot C,&b_t=M\cdot I_t \end{gather}$

である．適切に計算すれば， $(13)$ は $O ( N (\log Q + \log P)$ で求めることが出来る．

これだと道のりがかなり長く，あまり綺麗ではありません．そこで，クエリの解釈の仕方を変えてみます．

$A_t=x$ のとき， $\mathrm{chmax}(A_t,v)$ によって $A_t=v$ となるのは， $v=x,\ldots,M-1$ のときです．一方， $\mathrm{chmin}(A_t,v)$ によって $A_t=v$ となるのは， $v=0,\ldots,x-1$ のときです．そして，その他の場合では $A_t$ は変化しません ( $v=x$ のケースは都合のいいように解釈をしてもよいです)．ここで重要なのは， $x$ の値によらず， $A_t=0,\ldots,M-1$ へと変化するのが $1$ 通りずつあり，何も変化しないのが $M$ 通りあるということです．

つまり， $\mathrm{chmax}$ および $\mathrm{chmin}$ をひとまとめにして考えると，これらは次の 2 つの操作に置き換えられます．

$\mathrm{update}(l,r,v)$ : $A_l=\cdots=A_r=v$ へと更新する
$\mathrm{nop}(l,r,v)$ : 何もしない

ここまで来れば， $i$ 回の操作後に $A_t=x$ となる場合の数 $Y_t[i][x]$ を一発で求められます．まず， $A_t$ が $\mathrm{update}$ クエリの対象となるのは $b_t:=I_t\cdot M$ 通りです．また，クエリのパターン数は $a:=(2M+1)\cdot C$ 通りだったので， $\mathrm{update}$ クエリの対象とならないのは $a-b_t$ 通りということになります．

従って， $i$ 回の操作において一回も $\mathrm{update}$ クエリの対象とならないのは $(a - b _ t) ^ i$ 通りです．このとき，操作後の $A_t$ の値は $0$ です．

そして， $i$ 回の操作において一回以上 $\mathrm{update}$ クエリの対象となるのは $a ^ i - (a - b _ t) ^ i$ 通りです．このとき，操作後に $A_t=v\;(v=0,\ldots,M-1)$ となる場合の数は $v$ によりません．つまり， $A_t=v$ となる場合の数は $\dfrac{1}{M}\left(a ^ i - (a - b _ t) ^ i\right)$ 通りです．

以上より， $Y_t[i][x]$ は以下のように求められ，これは上で示した結果と一致しています．

$\displaystyle Y_t[i][x]=\begin{cases} \dfrac{1}{M}\left(a^i-(a-b_t)^i\right)+(a-b_t)^i & \text{(if $x=0$)}\\ \dfrac{1}{M}\left(a^i-(a-b_t)^i\right)& \text{(otherwise)} \end{cases}$

そして，任意の $i$ 回の操作における $\mathrm{sum}$ クエリの出力の総和 $S_t[i]$ に関する漸化式 :

$\displaystyle S_t[i+1]=a\cdot S_t[i]+I_t\sum_{x=0}^{M-1}x\cdot Y_t[i][x]$

に代入すると解ける形になるので，解いて $\displaystyle\sum_{t=0} ^ {N-1} S_t[Q]$ を計算すれば AC が取れます (代入と漸化式の計算パートは折り畳み部分に書いています)．

実装 (Python)

提出リンク (AC)

P = 998244353

N, M, Q = map(int, input().split())

inv_2 = pow(2, -1, P)
C = (N + 1) * N * inv_2 % P
I = [(t + 1) * (N - t) % P for t in range(N)]
a = (2 * M + 1) * C % P
b = [M * I[t] % P for t in range(N)]

ans = 0
for t in range(N):
    v1 = I[t] * (M - 1) % P * inv_2 % P
    u1 = Q * pow(a, Q - 1, P) % P
    u2 = (pow(a - b[t], Q, P) - pow(a, Q, P)) * pow(b[t], -1, P) % P
    ans += v1 * (u1 + u2) % P
print(ans % P)