Rust の競プロで同じ値を含む配列の順列を扱う方法 - ABC363-C を例に

はじめに

2024-07-20 の AtCoder ABC363-C で、同じ値を含む配列の順列を扱う問題が出題されました。

同じ値を含まない順列を考えます。 abc の 3文字を並び替えると次の 6通りです。 $_3P_3=3!=6$ 通り調べれば良いです。

abc
acb
bac
bca
cab
cba

同じ値を含む場合は、重複を除く必要があります。 aab の 3文字を並び替えると次の 3通りです。

aab
aba
baa

競技プログラミングの本番中、高速なはずの Rust で、通ったものの制限時間ギリギリでした。こうすれば実行時間が短くできるというのを、本番後に調べてみました。その記事です。

問題

詳しくは公式問題ページをどうぞ。

問題文

英小文字のみからなる長さ $N$ の文字列 $S$ が与えられます。
$S$ の文字を並び替えて得られる文字列（$S$ 自身を含む）であって、長さ $K$ の回文を部分文字列として 含まない ものの個数を求めてください。

制約

• $2≤K≤N≤10$
• $N,K$ は整数
• $S$ は英小文字のみからなる長さ $N$ の文字列

たとえば S=aab, K=2 のときは、正解 1通りです。

文字列	チェック
aab	aa が回文
aba
baa	aa が回文

本記事で扱うこと、要点

• 順列列挙ライブラリーが同じ値を含む配列に対してどうふるまうかはライブラリー次第
  • itertools::permutations() のように aab aab と 2回出力するものもあれば
  • permutohedron::next_permutation() のように aab 1回だけ出力するものもある
• 出力に同じ配列が含まれる場合、はじくためには、コレクションを使って登録済みかをチェックするのがよい
  • コレクションの種類: FxHashSet ≥ HashSet ≥ BTreeSet
  • コレクションに入れる値: usize > String ≥ Vec<char>
• 出力に同じ配列が含まれない場合は、はじく処理を省略できて簡単
• おまけ: この問題は abcdefghij が $10!$ 通りのように、1回だけ現れる文字を特別扱いすると高速に解ける

計測について

cargo build --release でビルドしたモジュールを、Surface Laptop 2 (Core i5-8250, 8GB) でそれぞれ5回実行しました。最も遅い時間を除いた 4回の平均値を比べます。

コードと計測結果です。記事中の表の値はここから取り出しました。

• AtCoder ABC363 C - Avoid K Palindrome 2 複数の実装とパフォーマンス測定
• AtCoder ABC363 C - Avoid K Palindrome 2 複数の実装とパフォーマンス測定 結果

実行速度はメモリスワップ等の影響を受けて大きく変わります。本記事中で実行時間の差が 2, 3割程度なら、状況によっては逆転するかもしれません。

重複ありの順列から、以前現れた組を除く方針

• Itertools in itertools - Rust > permutations()

Rust で競技プログラミングするときには、itertools を良く使うと思います。itertools に順列列挙機能があります。有力な選択肢になりそうです。

しかし、itertools::permutations() は同じ値の順列を考慮しません。ドキュメントに注意が書かれています。

Note: Permutations does not take into account the equality of the iterated values.

use itertools::Itertools;

let it = vec![2, 2].into_iter().permutations(2);
itertools::assert_equal(it, vec![
    vec![2, 2], // Note: these are the same
    vec![2, 2], // Note: these are the same
]);

そのため、itertools::permutations() を使うなら、同じ値をはじくように map などを使うことになります。

itertools::permutations() で解く流れ

// s が回文か調べる
fn is_palindrome(s: &[char]) -> bool {
    let k = s.len();
    (0..(k / 2)).all(|i| s[i] == s[k - i - 1])
}

// s を並び替えて得られる文字列のうち、長さ k の回文を含まない個数を返す 
fn f(k: usize, s: &[char]) -> usize {
    let n = s.len();
    let mut result = 0usize;
    for v in s.iter().cloned().permutations(n) {
        // ★重複チェックする

        if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
            result += 1;
        }
    }
    result
}

「// ★重複チェックする」について、これからどういう実装が効率的かを考えていきます。

k = 3, s = ['a', 'a', 'b'] で重複チェックをしない場合は、順列と回文チェック結果は次のようになります。

v: Vec<char>	回文チェック
['a', 'a', 'b']	['a', 'a'] が回文
['a', 'b', 'a']
['a', 'a', 'b']	['a', 'a'] が回文
['a', 'b', 'a']
['b', 'a', 'a']	['a', 'a'] が回文
['b', 'a', 'a']	['a', 'a'] が回文

同じ文字列を 2回カウントして、答えを 2としてしまっています。

文字の配列を set に入れて重複チェックする

• HashSet in std::collections - Rust

set で配列をそのまま調べます。 まずは基本の HashSet にします。Rust 以外ではハッシュを使っていれば間違いないことが多いです。

 fn f(k: usize, s: &[char]) -> usize {
     let n = s.len();
     let mut result = 0usize;
+    let mut set = HashSet::new();
     for v in s.iter().cloned().permutations(n) {
+        if set.contains(&v) {
+            continue;
+        }
+        set.insert(v.clone());
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
     }
     result
 }

これなら、重複チェックと回文チェックの両方をパスした 1つが答え、と正しい答えが得られます。

v: Vec<char>	重複チェック	回文チェック
['a', 'a', 'b']		['a', 'a'] が回文
['a', 'b', 'a']
['a', 'a', 'b']	出現済み
['a', 'b', 'a']	出現済み
['b', 'a', 'a']		['a', 'a'] が回文
['b', 'a', 'a']	出現済み

• よくある質問 · プログラミング言語Rust

RustのHashMapが遅いのはなぜですか？

しかし Rust の HashMap, HashSet は遅いです。FAQ に取り上げられる程度には。そこで、FxHashSet, BTreeSet についても試してみます。

• fxhash - Rust (衝突しやすい代わりに速いハッシュ)
• BTreeSet in std::collections - Rust (二分木, ハッシュより一般的に遅い)

-    let mut set = HashSet::new();
+    let mut set = FxHashSet::default();

-    let mut set = HashSet::new();
+    let mut set = BTreeSet::new();

k = 5, s = "abcwxyzyxw" の計測結果です:

重複チェック方法	実行時間 [s]
HashSet<Vec<char>>	1.8305
FxHashSet<Vec<char>>	1.1071
BTreeSet<Vec<char>>	1.8130

FxHashSet を使っておけば良さそうです。

文字列を set に入れて重複チェックする

文字の配列よりも文字列の方が効率的かも？ と調べてみました。

 fn f(k: usize, s: &[char]) -> usize {
     let n = s.len();
     let mut result = 0usize;
     let mut set = HashSet::new();
     for v in s.iter().cloned().permutations(n) {
-        if set.contains(&v) {
+        let string: String = v.iter().collect();
+        if !set.insert(string) {
             continue;
         }
-        set.insert(v.clone());
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
     }
     result
 }

k = 5, s = "abcwxyzyxw" の計測結果です:

重複チェック方法	実行時間 [s]
HashSet<String>	1.4690 (変更前: 1.8305)
FxHashSet<String>	1.4059 (変更前: 1.1071)
BTreeSet<String>	1.9955 (変更前: 1.8130)

Vec<char> を String に切り替えても、あまり速度は変わりません。

数値を set に入れて重複チェックする

この問題では 'a'～'z' の 26種類の文字が、最大 10文字現れます。状態の種類は $26^{10} < 32^{10} = 2^{5 \times 10}$ です。64bit 整数型に乗ります。

fn chars_to_key(s: &[char]) -> usize {
    s.iter().fold(0, |acc, &c| acc * 32 + char_to_key(c))
}

 fn f(k: usize, s: &[char]) -> usize {
     let n = s.len();
     let mut result = 0usize;
     let mut set = HashSet::new();
     for v in s.iter().cloned().permutations(n) {
-        if set.contains(&v) {
+        let key = chars_to_key(&v);
+        if !set.insert(key) {
             continue;
         }
-        set.insert(v.clone());
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
     }
     result
 }

v: Vec<char>	key	重複チェック	回文チェック
['a', 'a', 'b']	1		['a', 'a'] が回文
['a', 'b', 'a']	32
['a', 'a', 'b']	1	出現済み
['a', 'b', 'a']	32	出現済み
['b', 'a', 'a']	1024		['a', 'a'] が回文
['b', 'a', 'a']	1024	出現済み

最大長さ10 の文字列の一致を調べるのは大変ですが、数値 1つなら簡単に比べられます。

k = 5, s = "abcwxyzyxw" の計測結果です:

重複チェック方法	実行時間 [s]
HashSet<usize>	0.8945 (変更前: 1.8305)
FxHashSet<usize>	1.3353 (変更前: 1.1071)
BTreeSet<usize>	0.9462 (変更前: 1.8130 )

HashSet, BTreeSet が、かなり速くなりました。 FxHashSet は逆に遅くなっています。この問題については FxHashSet が銀の弾丸にはならなそうです。

重複のない最悪ケースでの比較

set に値を入れて重複を調べる方針の場合、重複がないと set ではじくことがありません。ただメモリと検索時間を消費することになります。すべて異なる 10 文字の場合、$10! = 3,628,800$ 通りの順列がすべて set に入ります。

この最悪ケースについて、HashSet, FxHashSet, BTreeSet の性能を比べてみます。

k = 5, s = "abcdefgxyz" の計測結果です:

重複チェック方法	実行時間 [s]
なし	0.5711
HashSet<Vec<char>>	4.3600
FxHashSet<Vec<char>>	3.7412
BTreeSet<Vec<char>>	4.4947
HashSet<String>	3.3093
FxHashSet<String>	3.3425
BTreeSet<String>	3.5889
HashSet<usize>	1.2219
FxHashSet<usize>	1.7746
BTreeSet<usize>	1.3861

• 重複チェックをするなら、判定が軽い usize などを使いたい
• コレクションの種類は一長一短あるけれども、速度差は倍もなく、制限時間以内になるならどれを使っても良い

ということが分かります。

permutohedron::heap_recursive() で解く流れ

• permutohedron::heap_recursive - Rust

AtCoder で使える permutohedron の利用例です。

itertools::permutations() と似たコードです。要素をいくつ取り出すかと言うオプションはありません。

 fn f(k: usize, s: &[char]) -> usize {
     let n = s.len();
     let mut result = 0usize;
     let mut set = HashSet::new();
-    for v in s.iter().cloned().permutations(n) {
+    permutohedron::heap_recursive(&mut s, |s| {
        let key = chars_to_key(&v);
        if !set.insert(key) {
             continue;
         }
         set.insert(v.clone());
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
-    }
+    });
     result
 }

FxHashSet<usize> で実行時間を比べると次のようになりました。

順列列挙方法	5, abcwxyzyxw 実行時間 [s]	5, abcdefgxyz 実行時間 [s]
itertools::permutations()	1.3121	1.7746
permutohedron::heap_recursive()	1.4141	1.8762

だいたい同じです。

重複のない順列を使う

重複のない順列列挙をするライブラリーもあります。

permutohedron::next_permutation()

• permutohedron::LexicalPermutation - Rust > next_permutation()

たとえば permutohedron を使って let mut v = ['a', 'a', 'b']; に対して v.next_permutation() すると、v が次の順列 ['a', 'b', 'a'] に更新されます。

同じ ['a', 'a', 'b'] にはなりません。先ほどの itertools::permutations() や permutohedron::heap_recursive() のように状態を持つ場所がありません。 ['a', 'a', 'b'] の次がもし ['a', 'a', 'b'] だとすると、いつまでたってもその繰り返しから抜けられなくなります。

v: Vec<char>	回文チェック	v.next_permutation() 後の v
['a', 'a', 'b']	['a', 'a'] が回文	['a', 'b', 'a']
['a', 'b', 'a']		['b', 'a', 'a']
['b', 'a', 'a']	['a', 'a'] が回文	-

重い処理だった重複チェックを取り除けますので、とても高速化が期待できます。

はじめに s を並び替え、s.next_permutation() 次の順番に更新して true を返す間ずっとループを回します。

fn f(k: usize, s: &[char]) -> usize {
    use permutohedron::LexicalPermutation;

    let n = s.len();
    let mut s: Vec<_> = s.iter().cloned().sorted().collect();
    let mut result = 0usize;
    loop {
        if (0..=(n - k)).all(|i| !is_palindrome(&s[i..(i + k)])) {
            result += 1;
        }
        if !s.next_permutation() {
            break;
        }
    }
    result
}

前のコードと比較するとこんな感じです。 set がなくなりました。

 fn f(k: usize, s: &[char]) -> usize {
+    use permutohedron::LexicalPermutation;
+
     let n = s.len();
+    let mut s: Vec<_> = s.iter().cloned().sorted().collect();
     let mut result = 0usize;
-    let mut set = HashSet::new();
-    for v in s.iter().cloned().permutations(n) {
-       let key = chars_to_key(&v);
-       if !set.insert(key) {
-            continue;
-        }
-        set.insert(v.clone());
+    loop {
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
+        if !s.next_permutation() {
+            break;
+        }
     }
     result
 }

superslice::next_permutation()

• superslice::Ext - Rust > next_permutation()

AtCoder で使える superslice の利用例です。ほとんど変わりません。

 fn f(k: usize, s: &[char]) -> usize {
-    use permutohedron::LexicalPermutation;
+    use superslice::*;

     let n = s.len();
     let mut s: Vec<_> = s.iter().cloned().sorted().collect();
     let mut result = 0usize;
     loop {
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
         if !s.next_permutation() {
             break;
         }
     }
     result
 }

ライブラリーなしで順列列挙する

次のようなコードで順列列挙できます。

 fn f(k: usize, s: &[char]) -> usize {
-    use permutohedron::LexicalPermutation;
-
     let n = s.len();
     let mut s: Vec<_> = s.iter().cloned().sorted().collect();
     let mut result = 0usize;
     loop {
         if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
             result += 1;
         }
-        if !s.next_permutation() {
+        let Some((i, _)) = s.windows(2).enumerate().rev().find(|(_, v)| v[0] < v[1]) else {
             break;
         }
+        let Some((j, _)) = s.iter().enumerate().rev().find(|(_, &c)| s[i] < c) else {
+            unreachable!();
+        };
+        s.swap(i, j);
+        s[(i + 1)..].reverse();
     }
     result
 }

aabc の並び替えイメージです。この方法で昇順に $3! \times 2 = 12$ 通りの順列を得られます。

配列	v[i] < v[i+1] i最大	v[i] < v[j] j最大	v.swap(i,j)	v[i+1].reverse()
aabc	2 bc	3 c	aac b	aac b
aacb	1 ac	3 b	ab ca	ab ac
abac	2 ac	3 c	abc a	abc a
abca	1 bc	2 c	ac ba	ac ab
acab	2 ab	3 b	acb a	acb a
acba	0 ac	2 b	b caa	b aac
baac	2 ac	3 c	bac a	bac a
baca	1 ac	2 c	bc aa	bc aa
bcaa	0 bc	1 c	c baa	c aab
caab	2 ab	3 b	cab a	cab a
caba	1 ab	2 b	cb aa	cb aa
cbaa	-	-	-	-

速度比較

順列列挙方法	5, abcwxyzyxw 実行時間 [s]	5, abcdefgxyz 実行時間 [s]
itertools::permutations() + 重複をはじく	0.8649	1.2219
permutohedron::next_permutation()	0.0157	0.1000
superslice::next_permutation()	0.0148	0.0986
ライブラリーなし	0.0132	0.0939

速度が 1桁違います。実装が簡単ですし、良いことばかりです。

おまけ: さらに高速化する

0.1s を切っていれば、ここまでで問題を解くには十分です。Rust ほど速くない言語で書いたとしても、十分制限時間に間に合いそうです。

せっかくですので、さらに高速化してみます。

1回しか現れない文字を特別扱いする

重複無しの順列を使うと十分な速度が出ますが、それでも $10!$ 通りのループを回しています。もっと減らしたいです。

ここで、abcdefghij と abcdefgxyz を考えます。すべての文字が 1回ずつ現れるということでは変わらないので、同じ答えになるはずです。

1回ずつ現れるものを「同じ文字、ただし回文にはならない」という扱いにしてみます。そうすると、調べる数が激減します。

abcdefghij は aaaaaaaaaa の順列 1通りに、$10!$ をかければ終了します。 aabbccddee の $10! / 2^5$ 通りが一番重い問題になります。

1回だけ現れる文字を a に寄せる関数と、回文チェックで a を対象外にする関数を用意します。

fn is_same_ex(i0: char, i1: char) -> bool {
    i0 == i1 && i0 != 'a'
}

fn is_palindrome_ex(s: &[char]) -> bool {
    let k = s.len();
    (0..(k / 2)).all(|i| is_same_ex(s[i], s[k - i - 1]))
}

// 文字列を 1回だけ現れるものをすべて 'a' に、それ以外を 'b' から始まる順に組み立て直す
fn build_s_ex(s: &[char]) -> (Vec<char>, usize) {
    let n = s.len();
    let mut a = [0usize; 26];
    for &c in s {
        let i = char_to_key(c);
        a[i] += 1;
    }

    let mut counts = vec![0usize; n];
    for k in a {
        if k > 0 {
            counts[k - 1] += 1;
        }
    }
    let n0 = counts[0];

    let mut s = vec!['a'; n0];
    let mut cur = 'b';
    for (i, &k) in counts[1..].iter().enumerate() {
        for _ in 0..k {
            s.append(&mut vec![cur; i + 2]);
            cur = (cur as u8 + 1) as char;
        }
    }

    // next_permutations() ですべて辿るために昇順の必要あり。そう組み立て済みのはず
    s.sort();

    // 1回だけ現れる文字はどこに現れても良いので、最後に階乗をかけられるようにする
    let m = (1..=n0).fold(1, |acc, x| acc * x);

    (s, m)
}

superslice::next_permutation() のサンプルを書き換えてみます。

 fn f(k: usize, s: &[char]) -> usize {
     use superslice::*;

     let n = s.len();
-    let mut s: Vec<_> = s.iter().cloned().sorted().collect();
+    let (mut s, m) = build_s_ex(s);
     let mut result = 0usize;
     loop {
-        if (0..=(n - k)).all(|i| !is_palindrome(&v[i..(i + k)])) {
+        if (0..=(n - k)).all(|i| !is_palindrome_ex(&s[i..(i + k)])) {
             result += 1;
         }
         if !s.next_permutation() {
             break;
         }
     }
-    result
+    result * m
}

superslice::next_permutation() の測定結果です

1回出現	5, abcwxyzyxw 実行時間 [s]	5, abcdefgxyz 実行時間 [s]	2, aabbccddee 実行時間 [s]	10, aabbccddee 実行時間 [s]
対応なし	0.0148	0.0986	0.0033	0.0018
対応あり	0.0010	0.0000	0.0035	0.0017

一番遅くても 3.5ms です。お疲れさまでした。

再帰ですぐに重複・回文をはじく

itertools で set で重複をはじくのが遅かったのは、組み立てられた長い文字列に対して確認しているためでした。1文字ずつ調べれば、重複はその階層だけで考えれば良いです。全然重くなりません。

k = 2, s = "aab"

1文字目	2文字目	3文字目	重複チェック	回文チェック
a
a	a
a	a	b		['a', 'a'] が回文
a	b
a	b	a
a	a		出現済み
b
b	a
b	a	a		['a', 'a'] が回文
b	a		出現済み

こうすると、最後の 1文字だけ重複判定すれば十分になります。かなり速くなります。

さらに、1文字ずつ調べるということは、2文字目時点で aa 回文が見つかれば 3文字目以降を調べる必要がなくなります。

1文字目	2文字目	3文字目	重複チェック	回文チェック
a
a	a			['a', 'a'] が回文
a	b
a	b	a
a	a		出現済み
b
b	a
b	a	a		['a', 'a'] が回文
b	a		出現済み

更に一つ前のように、1回だけ現れる文字をまとめて扱う処理を組み合わせられます。すべて行うと次のようになります。

fn f(k: usize, s: &[char]) -> usize {
    let (mut s, m) = build_s_ex(s);
    let result = f_recur(k, &mut s, 0);
    result * m
}

fn f_recur(k: usize, s: &mut [char], i: usize) -> usize {
    let n = s.len();

    let mut result = 0;
    let mut used = [false; 6]; // aabbccddee
    let palindrome_candidate = i + 1 >= k && is_palindrome_ex(&s[(i + 2 - k)..i]);
    for j in i..n {
        let key = char_to_key(s[j]);
        if used[key] {
            continue;
        }
        used[key] = true;
        s.swap(i, j);
        if !(palindrome_candidate && i + 1 >= k && is_same_ex(s[i], s[i + 1 - k])) {
            if i + 1 == n {
                result += 1;
            } else {
                result += f16_recur(k, s, i + 1);
            }
        }
        s.swap(i, j);
    }

    result
}

方法	5, abcwxyzyxw 実行時間 [s]	5, abcdefgxyz 実行時間 [s]	2, aabbccddee 実行時間 [s]	10, aabbccddee 実行時間 [s]
superslice	0.0148	0.0986	0.0033	0.0018
superslice 1文字特別	0.0010	0.0000	0.0035	0.0017
再帰	0.0170	0.1157	0.0040	0.0027
再帰 途中で回文	0.0111	0.0807	0.0012	0.0028
再帰 途中で回文 + 1文字特別	0.0008	0.0000	0.0014	0.0027

再帰でも同等程度に速くなりました。一番遅くて 3ms を切っています。

最後に

次の話をしました:

• 順列列挙ライブラリーが同じ値を含む配列に対してどうふるまうかはライブラリー次第
  • itertools::permutations() のように aab aab と 2回出力するものもあれば
  • permutohedron::next_permutation() のように aab 1回だけ出力するものもある
• 出力に同じ配列が含まれる場合、はじくためには、コレクションを使って登録済みかをチェックするのがよい
  • コレクションの種類: FxHashSet ≥ HashSet ≥ BTreeSet
  • コレクションに入れる値: usize > String ≥ Vec<char>
• 出力に同じ配列が含まれない場合は、はじく処理を省略できて簡単
• おまけ: この問題は abcdefghij が $10!$ 通りのように、1回だけ現れる文字を特別扱いすると高速に解ける

私は permutohedron や superslice の v.next_permutation() に対して、 v が書き変わるのは使い方が難しそうに思い、これまで避けていました。でも配列全体の clone をせず済むこと、この問題のように同じ文字が現れることを場合があるとすると、ありだなと思いました。

HashSet はやはり遅いです。 FxHashSet に切り替えても、この問題についてはそこまで効果がありませんでした。 $10!$ 個の文字列を set につめて重複チェックするというのは、あまり筋が良くないのでした。

また、印象的な問題があれば今回のように書いてみようかなと。