Magicode logo
Magicode
3 min read

RustとAtCoderを勉強する(typical90_b)

https://cdn.apollon.ai/media/notebox/da8782ef-954c-4923-bc3c-be5f45225299.jpeg

はじめに

AtCoder の問題を Rust で解いていきます。AtCoder も Rust も初心者ですが、温かい目で成長を見守っていただけるとありがたいです。

提出コード

途中 if 文が長くなってしまいました。もっとスマートに書く方法はないでしょうか...。
rust
use itertools::Itertools;
use proconio::input;

fn main() {
    input! {n: usize}

    let mut parentheses: Vec<String> = Vec::new();
    for bit in 0..(1 << n) {
        let s = (0..n)
            .map(|digit| if (bit & (1 << digit)) == 0 { '(' } else { ')' })
            .join("");
        if s.chars().filter(|c| *c == '(').count() == s.chars().filter(|c| *c == ')').count()
            && s.chars()
                .scan(0, |state, c| {
                    *state += if c == '(' { 1 } else { -1 };
                    Some(*state)
                })
                .all(|x| x >= 0)
        {
            parentheses.push(s);
        }
    }

    parentheses.sort();
    for s in parentheses {
        println!("{}", s);
    }
}

解説

長さ N のカッコ列のうち、正しいカッコ列(全ての括弧が対応している)を辞書順に列挙する問題です。今回の場合、N の最大値は 20 なので全パターン列挙しても 2^20 となり、全探索が可能です。提出コードのうち、カッコ列を全列挙する部分が以下のプログラムとなります。
rust
let n = 3;
for bit in 0..(1 << n) {
    let s = (0..n)
        .map(|digit| if (bit & (1 << digit)) == 0 { '(' } else { ')' })
        .collect::<String>();
    println!("{}", s);
}

(((
)((
()(
))(
(()
)()
())
)))
()
最後の出力は for 式の戻り値の () なので気にしないでください。ビットシフトビット論理積 を利用しています。
次に if 文にて、カッコ列が正しいカッコ列かどうかを判別しています。正しいカッコ列の条件は
  • 始まり('(')と終わり(')')が同じ数
  • カッコ列の最初から見ていったとき、「'(' の数」<「')' の数」となるケースが1回もない
となります。
最後に、正しいカッコ列を辞書順にソートして出力すれば解答となります。

まとめ

python では itertools.product('()', repeat=n) のようにすれば辞書順に列挙してくれるのですが、rust ではそのようなスマートな書き方はないのか気になります。

追記(2022/5/5)

itertools.product('()', repeat=n) の rust バージョンを考えてみました。multi_cartesian_product() を利用すると良さそうです。
rust
// itertools が import 出来ないのでここでは動作しません!
use itertools::Itertools;

let n = 3;
let s: Vec<String> = (0..n)
    .map(|_| "()".chars())
    .multi_cartesian_product()
    .map(|cs| cs.iter().collect())
    .collect_vec();

Discussion

コメントにはログインが必要です。