範囲for文でcombinationの列挙

これは何?

$n$ 個の要素から $k$ 個の要素を選び出す組み合わせは ${}_n\mathrm{C}_r=\frac{n!}{(n-k)!k!}$ 通りあります。 $n$ と $k$ を与えられたとき、そのような組み合わせをすべて列挙したい時があるので、C++で実装をしました。

Combinationの表現には、整数のbitを使います。下位n bitを要素に対応させて、選ばれた要素に対応するbitに1を立てることで組み合わせを表現します。例えば、 $n=5,k=2$ のときの列挙は次のようになってほしいです。

この記事ではさらに、列挙されたcombinationをコンテナとみなし、範囲for文を使って簡潔に列挙を記述できるような構造体を書きました。つまり、次のようなコードが書けます。

for (auto i : combinations(n, k)) {
    // iを使った処理を書く
}

範囲for文の仕様

cpprefjpによれば、C++17以降の範囲for文は、次のように通常のfor文に展開されます。

for ( for-range-declaration : for-range-initializer ) statement

は以下のように展開される:

{
  auto && __range = for-range-initializer;
  auto __begin = begin-expr;
  auto __end = end-expr;     // __begin と __end は異なる型でもよい
  for ( ; __begin != __end; ++__begin ) {
    for-range-declaration = *__begin;
    statement
  }
}

細かい所は置いておいて、どうやらコンテナのイテレータをbegin()で初期化して、end()と一致するまでインクリメントで通常のfor文を回しているようです。範囲for文で使いたいだけであれば、都合のいいイテレータを作ってしまえば、実際にコンテナが存在しなくてもよさそうです。

next_combination

辞書順で次に来るcombinationを返す関数です。整数のbitで表現した場合、combinationの辞書順は数値の昇順と同値です。
ネットで調べると、次のような記事が見つかります。
qiita.com
nemutage.hatenablog.jp
一つ目の記事では、次のような実装が紹介されています。

int next_combination(int sub) {
    int x = sub & -sub, y = sub + x;
    return (((sub & ~y) / x) >> 1) | y;
}

二つ目の記事では、一つ目の記事の実装が分かりやすく解説されています。私が書いたnext_combinationも基本的には同じ仕組みで動くので、詳しい解説はこちらを読んでください。

上の実装では、取得したbitをどうにか使いまわして並べ替えるために、除算が使われています。除算は速度が遅いので、除算を使わない実装を考えました。立っているbitの数を数えるpopcountを使うことで、天下り的な実装で除算を回避できます。

int next_combination(int sub) {
    int k = __builtin_popcount(sub);
    sub += sub & -sub;
    return sub + (1 << (k - __builtin_popcount(sub))) - 1;
}

ここでのkはわざわざ求めなくても定数なので、前計算をしておけば演算の数をさらに減らすことができます。

int next_combination(int sub, int k_bits) {
    sub += sub & -sub;
    return sub + (k_bits >> __builtin_popcount(sub));
}

ここで、k_bits = (1 << k) - 1です。

構造体にする

#include <assert.h>

struct combinations {
    using u64 = std::uint64_t;

    uint n;
    uint k;

    combinations(uint n_, uint k_) : n(n_), k(k_) { assert(n_ < 64 && n_ >= k_ && k_ > 0); }

    struct itr {
        u64 s;
        u64 t;

        bool operator!=(itr &right) { return s != right.s; }
        void operator++() {
            s += (s & -s);
            s += t >> __builtin_popcountll(s);
        }
        u64 operator*() { return s; }
    };

    itr begin() { return {(1ull << k) - 1, (1ull << k) - 1}; }
    itr end() { return {(1ull << n) | ((1ull << (k - 1)) - 1), (1ull << k) - 1}; }
};

冒頭で述べたような挙動を実現するには、コンテナと、そのイテレータを自作すればよいです。イテレータはcombinationそのものであるsと、1がkbit並んだ演算用の変数であるtしか保持していません。イテレータのインクリメントにnext_combinationを入れれば、combinationがたくさん入ったコンテナのイテレータを一つ進めたような挙動をしてくれます。他にはイテレータへのアクセスと比較、コンテナのbegin()とend()を直感通りに実装すればよいです。
私の実装したnext_combinationはk=0の場合をうまく処理できないので、n < 64 && n >= k && k > 0を満たさない引数は初めに弾いています。

速度について

AtCoderのコードテスト機能を使って簡易的な速度の測定を行いました。
$n=30,k=15$ として、以下の三つのコードの速度比較を行いました。
forループは ${}_{30}C_{15}=155117520$ 回回ります。

1. こちらのnext_combinationを使ったもの。
ビット演算 (bit 演算) の使い方を総特集！〜マスクビットから bit DP まで〜 #競技プログラミング - Qiita
実行時間:783ms

ソースコード

#include <iostream>
using namespace std;

int next_combination(int sub) {
    int x = sub & -sub, y = sub + x;
    return (((sub & ~y) / x) >> 1) | y;
}

int main() {
    int ans = 0;
    int n = 30, k = 15;
    for (int bit = (1 << k) - 1; bit < (1 << n); bit = next_combination(bit)) {
        ans++;
    }
    cout << ans << "\n";
}

2. 筆者のnext_combinationを使ったもの
実行時間:314ms

ソースコード

#include <iostream>
using namespace std;

int next_combination(int sub, int k_bits) {
    sub += sub & -sub;
    return sub + (k_bits >> __builtin_popcount(sub));
}

int main() {
    int ans = 0;
    int n = 30, k = 15;
    int k_bits = (1 << k) - 1;
    for (int bit = k_bits; bit < (1 << n); bit = next_combination(bit, k_bits)) {
        ans++;
    }
    cout << ans << "\n";
}

3. 構造体に成型したもの
実行時間:359ms

ソースコード

#include <assert.h>
#include <iostream>
using namespace std;

struct combinations {
    using u64 = std::uint64_t;

    uint n;
    uint k;

    combinations(uint n_, uint k_) : n(n_), k(k_) { assert(n_ < 64 && n_ >= k_ && k_ > 0); }

    struct itr {
        u64 s;
        u64 t;

        bool operator!=(itr &right) { return s != right.s; }
        void operator++() {
            s += (s & -s);
            s += t >> __builtin_popcountll(s);
        }
        u64 operator*() { return s; }
    };

    itr begin() { return {(1ull << k) - 1, (1ull << k) - 1}; }
    itr end() { return {(1ull << n) | ((1ull << (k - 1)) - 1), (1ull << k) - 1}; }
};

int main() {
    int ans = 0;
    for (auto i : combinations(30, 15)) {
        ans++;
    }
    cout << ans << "\n";
}

やはり除算を取り除いたことは速度に貢献しているようです。next_combinationを入れ替えると速度が2倍以上になりました。
構造体にしたものはそれより少し遅くなりましたが、どうせ使うときはforループの中の処理も走るため、許容範囲内だと思います。

最後に

筆者が競プロのライブラリのようなものを書いたのはこれが初めてです。はじめは遊びでいじってたんですが、思った以上によさそうなものができたので記事にしてしまいました。おかしい所や改善点など在りましたらご連絡ください。

追記(2024年4月27日)

範囲for文を用いた列挙の実装について、noshi91さんによる部分集合列挙のクリーンな実装を見かけたので、こちらを参考に記事内の実装を更新しました。

参考

cpprefjp.github.io
qiita.com
nemutage.hatenablog.jp
tech.garilog.com
noshi91.hatenablog.com