Rustで「すべての組み合わせ」を生成する

はじめに

競技プログラミングやパズル系の問題を解いていると、 「複数の配列から取りうる すべての組み合わせを列挙したい」という場面がよくあります。

Rust の標準ライブラリには直接のユーティリティはありませんが、 itertools クレートを使うとこの処理を非常に簡潔に書けます。

この記事では、直積生成に使える2つのメソッドを紹介します。

iproduct!

iproduct! は 複数のイテレータの直積（デカルト積）を生成するマクロです。2つでも3つでも、任意の数のイテレータを渡せます。タプルとして返ってくるので扱いやすく、ネストしたループを書くより読みやすいと思います。

use itertools::iproduct;

fn main() {
    let xs = vec![1, 2];
    let ys = vec!['a', 'b'];

    for (x, y) in iproduct!(xs, ys) {
        println!("{:?}", (x, y));
    }
}

出力例：

(1, 'a')
(1, 'b')
(2, 'a')
(2, 'b')

multi_cartesian_product

multi_cartesian_product は「複数のイテレータを動的にまとめて直積を取りたい」という場合に便利です。

use itertools::Itertools;

fn main() {
    let data = vec![
        vec![1, 2],
        vec![10, 20, 30],
        vec![100, 200],
    ];

    for combo in data.into_iter().multi_cartesian_product() {
        println!("{:?}", combo);
    }
}

出力例：

[1, 10, 100]
[1, 10, 200]
[1, 20, 100]
[1, 20, 200]
[1, 30, 100]
[1, 30, 200]
[2, 10, 100]
[2, 10, 200]
[2, 20, 100]
[2, 20, 200]
[2, 30, 100]
[2, 30, 200]

multi_cartesian_product のメリット

multi_cartesian_product の大きな利点のひとつは、イテレータとして扱えるため、前後に map や filter を自然に挟めることです。

iproduct! はマクロであり、直積生成の部分がコードの中心に来ますが、multi_cartesian_product はあくまでイテレータ変換の一段階として扱えるため、次のように処理の流れをそのままメソッドチェーンで書けるのが強みです。

use itertools::Itertools;

fn main() {
    let data = vec![
        vec![1, 2, 3],
        vec![11, 22, 33],
        vec![111, 222, 333],
    ];

    let result = data
        .into_iter()
        // 直積前の前処理
        .map(|v| v.into_iter().filter(|x| x % 2 == 1).collect_vec())
        // 直積
        .multi_cartesian_product()
        // 直積後の後処理
        .filter(|combo| combo.iter().sum::<i32>() >= 25);

    for value in result {
        println!("{:?}", value);
    }
}

実行環境

• rustc 1.89.0
• itertools 0.14.0