C++でnCkやnPkを全列挙する関数 - とある理系の備忘録

n!通りの順列を全列挙する関数はnext_permutationという備え付けの関数を使えばできたが、next_combinationなるものはどうやら存在しないようだった。
例えば「N個の要素からK個選んだ時の和」などを考えるとき、combinationのパターンを全列挙したくなる時が出てくる。
きっかけになったのが以下の問題。この問題だとNとKが固定されているのでこんな大げさに考えなくてもいいのだが、知っておいて損はないだろうと思った。
beta.atcoder.jp

関数の実装方法について考察している記事がいくつかあったので自分で一から作ろうかとも思ったが、実装がわかりやすく直接貼り付けていい感じに動くものが見つかったので今回はそれを使うことにした。時間があったらじっくり考えます。
参考にしたのはこのページ。
stackoverflow.com

以下がサンプルコード。NとKの値を入力するとN個からK個選ぶ組み合わせを全列挙する。

#include <bits/stdc++.h>
#define rep(i, m, n) for(int i = m; i < (n); i++)
#define print(x) cout << (x) << endl;
#define printa(x,n) for(int i = 0; i < n; i++){ cout << (x[i]) << " ";} cout << endl;
#define printa2(x,m,n) for(int i = 0; i < m; i++){ for(int j = 0; j < n; j++){ cout << x[i][j] << " ";} cout << endl;}
#define printp(x,n) for(int i = 0; i < n; i++){ cout << "(" << x[i].first << ", " << x[i].second << ") "; } cout << endl;
#define INF (1e18)
using namespace std;
typedef long long ll;
const ll MOD = 1e9 + 7;
typedef struct{
    int x;
    int y;
} P;
typedef struct{
    ll to;
    ll cost;
} edge;
typedef pair<ll, ll> lpair;
template <typename Iterator>
inline bool next_combination(const Iterator first, Iterator k, const Iterator last)
{
   /* Credits: Thomas Draper */
   if ((first == last) || (first == k) || (last == k))
      return false;
   Iterator itr1 = first;
   Iterator itr2 = last;
   ++itr1;
   if (last == itr1)
      return false;
   itr1 = last;
   --itr1;
   itr1 = k;
   --itr2;
   while (first != itr1)
   {
      if (*--itr1 < *itr2)
      {
         Iterator j = k;
         while (!(*itr1 < *j)) ++j;
         iter_swap(itr1,j);
         ++itr1;
         ++j;
         itr2 = k;
         rotate(itr1,j,last);
         while (last != j)
         {
            ++j;
            ++itr2;
         }
         rotate(k,itr2,last);
         return true;
      }
   }
   rotate(first,k,last);
   return false;
}

int main(){
    cin.tie(0);
    ios::sync_with_stdio(false);
    ll N,K;
    cin >> N >> K;
    vector<ll> v(N);
    iota(v.begin(), v.end(), 0);
    do{
        printa(v,K);

    }while(next_combination(v.begin(),v.begin() + K,  v.end()));
}

next_combination関数がメインとなる部分。この関数を通した後、先頭K要素を取得することで今回求める組み合わせを全て出力することができる。

なお、iotaは指定した数値から始まる数列を生成する関数。原始関数のι(イオタ)が由来だとか。
この関数を用いるとN!の全列挙ではTLEするようなケースでも組み合わせを求めることができる。
試しにいくつかの(N,K)でやってみたところ、(20,10)や(100,3)程度なら高速に動作した。
オーダーがどれくらいか見積もることができなかったけど、Nがそれほど大きくなければ $O(N^K)$ くらいで動いている気がする。
誰かちゃんとしたオーダーがわかる人がいたら教えてくだされ。

追記:
next_combinationで生成された、先頭からK個の数列がCombinationということは、その部分数列(ソート済み)をnext_permutationにかけてあげればpermutationの全列挙も可能だということに気づいた。
先ほどのコードのwhile文のところを

    vector<ll> v2(K);
    do{
        rep(i,0,K){
            v2[i] = v[i];
        }
        do {
            printa(v2,K);
        }while(next_permutation(v2.begin(), v2.end()));

    }while(next_combination(v.begin(),v.begin() + K,  v.end()));

みたいにすればv2にはnPkの全列挙が入るね。やったね。
計算時間はK!倍になるのかな、いや律速はnext_combinationだからそこまで変わらないのかな。よくわからず。

なんか全体的にもっと速い方法がある気がするなぁ。突き詰めてくと面白そう。