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\) と置きます。

  1. すると \(L\) のうち更新される最小のものは、\(R\) の要素の最大値よりも小さいもののうち最大のものです。これを \(L_i\) と置きます。(存在しない場合は、next shuffle が存在しません。)
  2. さらに \(L _ i\)\(R\) の要素のうち \(L_i\) より大きな最小のものと交換されます。これを \(R_j\) と置きます。
  3. あとは \(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_slicerotate_{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: