Rust の Itertools で個数可変の直積集合を列挙する

コード中では import や use 文（宣言？）を省略します。

ちょっとした困りごと

Rust の itertools には、順列、組み合わせなどを列挙する関数があり大変便利なんですが、個数を指定できる直積集合を列挙する関数は直接用意されていません。

Python の等価な関数

Python の itertools には、直積集合を列挙する product 関数があり、このように使えます。

for a, b, c in itertools.product(range(1, 5), repeat=3):
    print(a, b, c)

iproduct の弱点

Rust でも iproduct マクロを使うことで同様なことを実現できます。

fn main() {
    for (a, b, c) in iproduct!(1..5, 1..5, 1..5) {
        println!("{} {} {}", a, b, c);
    }
}

しかし、 iproduct マクロには以下の弱点があります。

• 多くの要素を必要とするとき、記述が長くなる
• イテレータを直接書くため、個数を可変にできない

fn main() {
    let size = 10; // この値を個数として使いたい
    
    // iproduct!(1..5, size) のように使いたい
    for x in iproduct!(1..5, 1..5, 1..5, 1..5, 1..5, 1..5, 1..5, 1..5, 1..5, 1..5) {
        // 列挙された直積の値の総和を計算し、出力する
        // タプルはイテレータにできないので x.iter().sum() が使えない!
        println!("{}", x.0 + x.1 + x.2 + x.3 + x.4 + x.5 + x.6 + x.7 + x.8 + x.9);
    }
}
        

解決策

そこで、現状存在するもので作ろうと思います。 multi_cartesian_product を使います。

fn main() {
    let size = 10; // 要素の個数
    
    for x in (0..size).map(|_| 0..5).multi_cartesian_product() {
        // 列挙された直積の値の総和を計算し、出力する
        println!("{}", x.iter().sum::<u64>());
    }
}
        

multi_cartesian_product は、イテレータ（にできるもの）のイテレータから、各要素の Vec からなる直積集合のイテレータを作ってくれます。今回は複数の集合としてすべて同じものを使いましたが、 map 関数の中身を作ることによって異なる集合の直積集合を作れます。

デメリット

そんな multi_cartesian_product ですが、デメリットも存在します。

• 要素の型がすべて同じでなければならない
• パターンマッチングによる代入ができない

ただ、上記の状況の場合 iproduct で事足りることが多く、さほど気にならないと思われます。[要出典]

使用例

これを使って、 ABC157-C を解きます。

fn solve(n: usize, sc: Vec<(usize, i64)>) -> i64 {
    // 判定関数
    let ok = |num: &Vec<i64>| {
        sc.iter().all(|&(s, c)| num[s] == c)
    };
    // コーナーケース: 0
    if n == 1 && ok(&vec![0; n]) {
        return 0;
    }
    // 先頭のみ 1..=9, それ以外 0..=9 の長さ n の直積集合
    for num in (0..n).map(|i| if i == 0 { 1..=9 } else { 0..=9 }).multi_cartesian_product() {
        if ok(&num) {
            return num.iter().fold(0, |acc, &x| acc * 10 + x);
        }
    }
    // 見つからない場合 -1 を返す
    -1
}

fn main() {
    // proconio::input で入力
    input! {
        (n, m): (usize, usize),
        sc: [(Usize1, i64); m],
    }
    println!("{}", solve(n, sc));
}

おわりに

ちょっとこの例微妙じゃないですかね...競プロにおいては、問題を直接解くよりも、ランダムテストで愚直解を書く用途のほうが使えそうです。

おまけ

bit全探索で使えそうです。普通に整数でループ回したほうが早く書けそうですが、 bit 演算のことを考えなくて済みます。

fn solve(n: usize, ab: &[(usize, usize)], k: usize, cd: &[(usize, usize)]) -> usize {
    let mut ans = 0;
    // k-bit の組み合わせを全列挙
    for choose in (0..k).map(|_| [true, false]).multi_cartesian_product() {
        // 各箱にボールが置かれているか管理する
        let mut placed = vec![false; n];
        for (i, &c) in choose.iter().enumerate() {
            if c {
                placed[cd[i].0] = true;
            } else {
                placed[cd[i].1] = true;
            }
        }
        // 満たす条件の数を数える
        let mut cnt = 0;
        for &(a, b) in ab {
            if placed[a] && placed[b] {
                cnt += 1;
            }
        }
        ans = ans.max(cnt);
    }
    ans
}

fn main() {
    input! {
        (n, m): (usize, usize),
        ab: [(Usize1, Usize1); m],
        k: usize,
        cd: [(Usize1, Usize1); k],
    }
    println!("{}", solve(n, &ab, k, &cd));
}