yukicoder No.1051 PQ Permutation

※当初の公開時点での解説およびリンク先コードが誤っていたため、修正しました。

解法

※以下の解説中で「辞書順で1つ進める」と書いた時には、「もし進められない（既に辞書順最大である）ときには -1 を出力して終了する」という意味を含むものとします。実装上はC++であれば next_permutation を用いると楽です。

答え $B$ は与えられた $A$ より辞書順で大きくないといけないので、まず与えられた $A$ を辞書順で1つ進めます。この状態で $P, Q$ の位置をチェックして、 $P$ が $Q$ より左にあればこれが答えです。

$Q$ のほうが左にある場合、先頭から $Q$ までの並びがこのままである限りは条件を絶対に満たせません。よって先頭から $Q$ までの並びが崩れるまで順列を進めます。具体的には $Q$ より後ろ（ $Q$ を含まない）を逆順ソートして、さらに辞書順で1つ進めます。

この操作後に順列がどうなるかを場合分けします。（-1 で終了するパターンは書いていません）

f:id:betrue12:20200509005123p:plain

細かく分けるとこんな感じですが、パターン1とパターン3、パターン2とパターン4はほぼ同じです。 $Q$ より後ろを逆順ソートしてさらに1つ進めると、 $Q$ もしくはそれより前の位置のどこかが「繰り上がり」、それ以降は余った値が昇順ソートされた順列になります。

これらのパターンにおける $P, Q$ の位置関係を見ると、

パターン1, 3のように「繰り上がり」に $P$ が使われた場合、必ず $P$ が $Q$ よりも左に来るので、これが答えとなる。
パターン2, 4のように「繰り上がり」に $P$ も $Q$ も使われなかった場合、 $P, Q$ は昇順ソートされた領域に属するため、その位置関係は大小に依存する。 $P \lt Q$ であればそれが答えであり、 $Q \lt P$ であれば $Q$ を $P$ のすぐ右隣に移動させるのが最適である。
パターン5のように「繰り上がり」に $Q$ が使われた場合、必ず $Q$ が $P$ よりも左に来る。

となります。パターン5以外の場合はこれで答えが求められます。

パターン5だけは決着が付かないので、「再度 $Q$ より後ろを逆順ソートする→さらに辞書順で1つ進める→繰り上がりパターンチェック→パターン5なら再度…」という流れを繰り返すことになります。これは最悪ケースで約 $N$ 回繰り返すことになるので、毎回ソートと next_permutation を実行していると間に合いません。しかし次のようにやると間に合わせることができます。

そもそも「さらに辞書順で1つ進める」処理の時に繰り上がる位置はどこなのか考えます。それは $Q$ が存在する位置から左に見ていって、初めて $A_{p} \lt A_{p+1}$ となるような位置 $p$ です。そして $A_{p+1}, ..., A_{N}$ の中で $A_{p}$ より大きいもののうち最小の値が位置 $p$ に来ることによって繰り上がります。

もし位置 $p$ に来る値が $Q$ 以外であれば、先述のパターン5から逃れられます。そして $Q$ が来てしまう場合は、「辞書順で1つ進めて、再び $Q$ より後ろを逆順ソート」という一連の処理の代わりに、 $A_{p}$ と $Q$ をスワップしてしまうことができます。

その理由は以下の通りです。繰り上がり処理を行う前の状態では $A_{p+1}, ..., A_{N}$ が降順ソートされていました。そして「 $A_{p+1}, ..., A_{N}$ の中で $A_{p}$ より大きいもののうち最小の値」が $Q$ だったので、値 $Q$ があった位置に値 $A_{p}$ が代わりに入っても位置 $p+1$ 以降における相対的な順序は変わりません。そのため「 $Q$ より後ろが降順ソートされている」という条件が保たれます。

f:id:betrue12:20200509025201p:plain

繰り上がりの位置 $p$ をずらしていくのは全体で $O(N)$ 、繰り上がりで位置 $p$ に来る値の判定は set 等を用いて1回あたり $O(\log N)$ 、値のスワップは1回あたり $O(1)$ です。よってこの処理を約 $N$ 回繰り返しても大丈夫です。 $Q$ が繰り上がりに使われるパターン5を抜け出して答えに辿り着くか、または $A_{p} \lt A_{p+1}$ となるような $p$ が存在しなくなってしまい -1 で終了するまで繰り返せば良いです。