ARC174-Eが面白かったので備忘録。https://atcoder.jp/contests/arc174/tasks/arc174_e

問題は、$${1 \le K \le N \le 3\times 10^5}$$ と 数列 $${P = \{P_1, …, P_N\}, 1 \le P_i \le N, P_i \ne P_j \text{ if } i \ne j}$$ が与えられ、また $${t = 1,…,N}$$ の全てについて、辞書式順序で$${P}$$ 以下の数列

$$
\lbrace a = \{a_1,…,a_N\} ~|~  1 \le a_i \le N, a_i \ne a_j \text{ if } i \ne j, t \in a \rbrace
$$

の個数を数えよ、という問題。

答えを $${N}$$ 個提出しないといけなくて、$${N \le 3\times 10^5}$$ なので、$${t}$$ を一個ずつ回している場合では無さそう。$${P_i}$$ について一回の処理を $${ \mathcal{O}(\log N)}$$ くらいで終わらせないといけなくて、それでいて $${10^5}$$ 個の回答を引き出すのだから、なにかまとめて処理できるに違いない。と思いつつ場合の数を数える。

一旦 $${t}$$ のことを考えなければ、$${P}$$ 以下の数列としてあり得るのは

• $${1 \le a_1 < P_1, \text{ any } a_2,…,a_N}$$

• $${a_1 = P_1, 1 \le a_2 < P_2, \text{ any } a_3,…,a_N}$$

• $${a_1 = P_1, a_2 = P_2, 1 \le a_3 < P_3, \text{ any } a_4,…,a_N}$$

• …

で、数を数えると

• $${a_1 < P_1}$$ の場合

  • $${ (P_1 - 1) {}_{N-1}P_{K-1} }$$ 個

• $${a_1 = P_1, a_2 < P_2}$$ の場合

  • $${P_1 < P_2}$$ なら $${ (P_2 -2){}_{N-2}P_{K-2} }$$ 個

  • $${P_2 < P_1}$$ なら $${ (P_2 -1){}_{N-2}P_{K-2} }$$ 個

• $${a_1 = P_1, a_2 = P_2, a_3 < P_3}$$ の場合

  • $${P_1, P_2 < P_3}$$ なら $${ (P_3 - 3){}_{N-3}P_{K-3} }$$ 個

  • $${P_1 < P_3 < P_2}$$ または $${P_2 < P_3 < P_1}$$ なら $${ (P_3 -2){}_{N-3}P_{K-3} }$$ 個

  • $${P_3 < P_1, P_2}$$ なら $${ (P_3 -1){}_{N-3}P_{K-3} }$$ 個

• …

と、$${P_i}$$ までに既に使った文字のなかで $${P_i}$$ より小さいものが何個あるかによってカウントする個数が違いそう。

そこで、この順で加算することを想定し、$${ a_1=P_1,…,a_{i-1}=P_{i-1}}$$ まで決まっていて、次は $${a_i < P_i}$$ となるような選び方が何個あるかを数えることにする。

$${P_i}$$ より小さくてまだ$${P_1,…,P_{i-1}}$$に出ていない数の個数を$${c_i}$$とすると、$${a_i}$$ として選べるのは $${c_i}$$ 個。また、$${a_i < P_i}$$ なので残りは任意でとってよいので、残りの数字の個数及び残りの枠数をそれぞれ$${n = N - (i - 1), k = K - (i - 1)}$$ とすれば：

$${t}$$ が既に $${P_1,…,P_{i-1}}$$ の中で使われていれば、あとは $${t}$$ を気にせず自由に決めていいので、$${a_i}$$ は $${c_i}$$ 通り、それ以降は $${n-1}$$個の中から順番に $${k - 1}$$ 個取り出せばよくて、$${{}_{n-1}P_{k-1}}$$ 通り。合わせて、$${c_i \times {}_{n-1}P_{k-1}}$$ 通り。

$${t}$$ がまだ使われていなければ、次以降で使わなければならない。$${t < P_i}$$ か $${ P_i \le t}$$ かで $${t}$$ を$${a_i}$$にできるかどうかが違うので、場合の数も違うのでここで場合分けする。

$${t < P_i}$$ の場合はさらに $${a_i = t}$$ と $${a_i \ne t}$$ があり、前者は残りは任意でいいので $${{}_{n-1}P_{k-1}}$$通りと、後者は $${(c_i-1)(k-1){}_{n-2}P_{k-2}}$$ 通り。$${c_i-1}$$は$${t}$$以外で$${a_i}$$になれる個数で、$${k-1}$$ は$${t}$$の場所を決め、残りは任意。

他方、$${P_i \le t}$$ の場合は、$${a_i}$$は$${t}$$にできないので、$${c_i(k-1){}_{n-2}P_{k-2}}$$ 通り。$${c_i}$$ は$${a_i}$$の個数、$${k-1}$$ は$${t}$$の置く場所で、残りが$${{}_{n-2}P_{k-2}}$$。

まとめると、$${ a_1=P_1,…,a_{i-1}=P_{i-1}}$$ まで決まっていて、次は $${a_i < P_i}$$ であるような場合の数は、

• ① $${t \in \{P_1,…,P_{i-1}\}}$$ の場合

  • $${c_i \times {}_{n-1}P_{k-1} }$$ 通り

• ② $${t \not\in \{P_1,…,P_{i-1}\}}$$ で、

  • $${t < P_i}$$ の場合

    • $${ {}_{n-1}P_{k-1} + (c_i-1)(k-1){}_{n-2}P_{k-2} }$$ 通り

  • $${P_i \le t}$$ の場合

    • $${  c_i(k-1){}_{n-2}P_{k-2} }$$ 通り

これを、$${i=1,…,K+1}$$ の順に、各回で全$${t}$$ に対して加算していけばよい。（$${i=K+1}$$ は $${a_1=P_1,…,a_K=P_K}$$ が決まっている場合で、1通り。）

ところで、これを愚直に加算しようとすると、各$${i}$$ について $${t}$$ 個の場所があるので、$${\mathcal{O}(N^2)}$$ かかってしまってTLE。幸い、全ての加算は区間に同じ値をたすだけなので、これは区間加算で処理すれば良さそう。また、①か②かは $${i}$$ が進むにつれて高々一度だけ切り替わるので、②は全ての$${t}$$ についてずっと加算し続けながら、②から抜けて①になった時に $${t}$$ の分だけ取り出して保存しておけばよい（それ以降は②の側で加算される値は無視していい）。そこで、modint が区間加算、一点取得ができる LST を使えばOK。$${P_i}$$ が来た時点で $${t=P_i}$$ の場所の値を取り出せば、それが②側の$${t}$$ での値になっている。

①の側は $${t \in \{P_1,..,P_{i-1}\}}$$ なる $${t}$$ について各回で同じ値をそれまでの $${t \in \{P_1,..,P_{i-1}\}}$$ に加算すればよくて、これは並べ替えると

• $${t=P_1}$$ は、$${ \sum_{i=1}^K c_i\times {}_{N-i}P_{K-i} }$$

• $${t=P_2}$$ は、$${ \sum_{i=2}^K c_i\times {}_{N-i}P_{K-i} }$$

• $${t=P_3}$$ は、$${ \sum_{i=3}^K c_i\times {}_{N-i}P_{K-i} }$$

• …

これは各回の$${c_i\times {}_{N-i}P_{K-i}}$$ を保存しておいてあとで累積和を取ればOK。（今回はLSTを使ったけれど。）

このように、各 $${t}$$ について①と②を別々に計算しておいて、あとで合算すればめでたく出力する値が求まる。