組み合わせを列挙するイテレータ

M種類のボールの中からN個以下選びだす場合 (0個の場合も1通りと数える) の,全ての場合についてループしたいとする.
ちなみにこういう問題の全部の通り数はcomb(M + N, M).つまりM+N個のボールからN個のボールを取り出すときの通り数と一致することが知られている. (ここではN個*以下*でない)

最初Haskellで書いて楽勝だったのだけど,C++で書こうとすると再帰とループがごちゃごちゃになるのでかなり面倒だ.
もちろん,再帰を使ってリストを出した後,リストにそってイテレーションする方法もあるのだけど,たいていこういう組み合わせの数は大きくなるのであまりオススメできない.
なので,結局自前のスタックを使って書くことにした.

struct CombIterator {
  size_t nvars;
  size_t maxorder;
  size_t curindex;

  std::vector<size_t> varstack;

  size_t maxidx;
  
  long comb(long l, long r) {
    unsigned long ret = 1;
    unsigned long den = 2;
    for (unsigned long m = 0; m < r; ++ m) {
      ret *= (l - m);
      if (den <= r && (ret % den) == 0) ret /= (den++);
    }
    return ret;
  }

  CombIterator(size_t nv, size_t mo)
    : nvars(nv), maxorder(mo), curindex(0) {
    maxidx = comb(nv + mo, mo);
  }

  size_t top() { return varstack[varstack.size() - 1]; }

  void next() {
    if (varstack.size() > maxorder) {
      return;
    }

    ++ curindex;

    size_t curorder = varstack.size();
    while (varstack.size() > 0 && top() == (nvars - 1)) {
      varstack.pop_back();
    }

    if (varstack.size() == 0) {
      for (size_t s = 0; s <= curorder; ++ s) varstack.push_back(0);      
    } else {
      size_t nv = top() + 1;
      varstack.pop_back();
      varstack.push_back(nv);
      while (varstack.size() < curorder) varstack.push_back(nv);
    }

  }

  bool done() { 
    if (varstack.size() > maxorder) {
      assert(maxidx == curindex);
      return true;
    } else {
      return false;
    }
  }

  size_t index() { return curindex; }
};

ボールの種類の数Nと取り出すボールの最大数Mの時,CombIterator(N, M)みたいな感じでインスタンス化して,done()がtrueになるまで,next()しつづけると,varstackの値がコロコロ変わってくれるという寸法だ.
なかなか便利.

ちなみにHaskellで書くと以下.

comblist1 :: [a] -> Int -> [[a]]
comblist1 _ 0 = [[]]
comblist1 [] _ = []
comblist1 allv@(v:vs) m = [v:rest | rest <- comblist1 allv (m - 1)] ++ comblist1 vs m
comblistN :: [a] -> Int -> [[a]]
comblistN vs n = foldr (++) [] (map (comblist1 vs) [0..n])

main = print$comblistN ['a', 'b', 'c', 'd', 'e', 'f', 'g'] 6

やはりエレガント!!