E - Modulo MST Editorial by ngtkana
next_permutation の要領で shuffle を列挙することで \(O((M - N)\binom{M}{N - 1})\) 時間で解くことができます。
解法
長さ \(N\) の数列 \(A\) の置換であって、前 \(K\) 項、後ろ \(N - K\) 項がそれぞれソートされているものを \((K, N - K)\)-shuffle と呼びます。これを列挙しましょう。
そのために、いま \(A\) が \((K, N - K)\)-shuffle であるとして、\(A\) より大きな最小の \(A\) の \((K, N - K)\)-shuffle を計算する方法を考えましょう。記述の簡単のため \(A\) の要素は相異なると仮定しますが、そうでない場合も全く同様です。 前 \(K\) 項を \(L\)、後ろ \(N - K\) 項を \(R\) と置きます。
- すると \(L\) のうち更新される最小のものは、\(R\) の要素の最大値よりも小さいもののうち最大のものです。これを \(L_i\) と置きます。(存在しない場合は、next shuffle が存在しません。)
- さらに \(L _ i\) は \(R\) の要素のうち \(L_i\) より大きな最小のものと交換されます。これを \(R_j\) と置きます。
- あとは \(L _ {i+1}, \dots, L _ {K - 1}\) と \(R _ {j + 1}, \dots, R _ {N - K - 1}\) を繋げて一つの配列と思い、これを \(R _ {j + 1}\) が先頭に来るように回転すれば答えになります。(これ説明難しいです。)
以上の操作は償却 \(O ( N - K )\) 時間で実行できます。
ステップ 3 は、要素のコピーができるのでしたら、キューなどを用いて実装してもよいです。
コピーをしたくない、もしくは O(1) extra memory で行いたいということでしたら、標準ライブラリの swap_slice
と rotate_{left,right}
を用いるとよいです。
fn next_shuffle<T: Ord>(a: &mut [T], k: usize) -> bool {
let n = a.len();
if n == k {
return false;
}
let (left, right) = a.split_at_mut(k);
let Some(mut i) = left.iter().rposition(|x| x < right.last().unwrap()) else {
return false;
};
let mut j = right.iter().position(|x| &left[i] < x).unwrap();
std::mem::swap(&mut left[i], &mut right[j]);
i += 1;
j += 1;
let swap_len = (k - i).min(n - k - j);
left[k - swap_len..].swap_with_slice(&mut right[j..j + swap_len]);
left[i..].rotate_left(k.saturating_sub(i + swap_len));
right[j..].rotate_right((n - k).saturating_sub(j + swap_len));
true
}
参考記事: next_permutation と同様に列挙できるもの一覧
提出
Rust (102 ms): https://atcoder.jp/contests/abc328/submissions/47515019
posted:
last update: