Alex's Anthology of Algorithms Common Code for Contests in Concise C++
Mathematics / Combinatorics

6.2.4 Enumerating Combinations

6-Mathematics/6.2.4_Enumerating_Combinations.cpp

A combination is a subset of size $k$ chosen from a total of $n$ (not necessarily distinct) elements, where order does not matter.

The lexicographic successor advances the rightmost chosen position that can still move right, then packs every later position as far left as possible. Ranking and unranking use the combinatorial number system: at each position, count how many combinations would be skipped by choosing a smaller next value, then either add that count to the rank or subtract it while searching for the requested rank. Bitmask successors use the same order as increasing integers with a fixed popcount.

  • next_combination(lo, mid, hi, comp = std::less<>()) takes random-access iterators lo, mid, and hi as a range $[{\htmlClass{math-inline-code}{\texttt{lo}}}, {\htmlClass{math-inline-code}{\texttt{hi}}})$ of $n$ elements for which the function will rearrange such that the $k$ elements in $[{\htmlClass{math-inline-code}{\texttt{lo}}}, {\htmlClass{math-inline-code}{\texttt{mid}}})$ become the next lexicographically greater combination. The function returns true if such a combination exists, or false if $[{\htmlClass{math-inline-code}{\texttt{lo}}}, {\htmlClass{math-inline-code}{\texttt{mid}}})$ already consists of the lexicographically greatest combination of the elements in $[{\htmlClass{math-inline-code}{\texttt{lo}}}, {\htmlClass{math-inline-code}{\texttt{hi}}})$, in which case the range is reset to its first combination. The comparator comp defines the element ordering.
  • next_combination(n, a) rearranges a to become the next lexicographically greater combination of distinct integers in the range $[0, {\htmlClass{math-inline-code}{\texttt{n}}})$. The vector a must be sorted and contain distinct integers in the range $[0, {\htmlClass{math-inline-code}{\texttt{n}}})$.
  • next_combination_mask(x) interprets the bits of an integer x as a mask with 1-bits specifying the chosen items for a combination and returns the mask of the next lexicographically greater combination (that is, the lowest integer greater than x with the same number of 1 bits). Note that this does not generate combinations in the same order as next_combination(), nor does it work if the corresponding $n$ items are not distinct (in that case, duplicate combinations will be generated). It returns $0$ if no successor fits in uint64_t.
  • combination_by_rank(n, k, r) returns the combination of $k$ distinct integers in the range $[0, {\htmlClass{math-inline-code}{\texttt{n}}})$ that is lexicographically ranked $r$, where $r$ is a 0-based rank in the range $[0, \binom{n}{k})$.
  • rank_by_combination(n, a) returns an integer representing the 0-based rank of combination a, which must contain sorted distinct integers in $[0, {\htmlClass{math-inline-code}{\texttt{n}}})$.
  • next_combination_with_repeats(n, a) rearranges a to become the next lexicographically greater combination of not necessarily distinct integers in the range $[0, {\htmlClass{math-inline-code}{\texttt{n}}})$. The vector a must be sorted. Note that there is a total of $n \mathbin{\text{multichoose}} k$ combinations if repetition is allowed, where $n \mathbin{\text{multichoose}} k = \binom{n + k - 1}{k}$.

Overflow warning: Exact counts and ranks used by the ranking operations must fit in int64_t.

Time Complexity
  • $O(n)$ per call to next_combination(lo, mid, hi), where $n$ is the distance between lo and hi.
  • $O(k)$ per call to next_combination(n, a) and next_combination_with_repeats(n, a), where $k$ is the size of a.
  • $O(1)$ per call to next_combination_mask().
  • $O(n \cdot k)$ per call to combination_by_rank() and rank_by_combination().
Space Complexity
  • $O(k)$ auxiliary for combination_by_rank(n, k, r).
  • $O(1)$ auxiliary for all other operations.

Implementation

#include <algorithm>
#include <cstdint>
#include <functional>
#include <iterator>
#include <numeric>
#include <utility>
#include <vector>

template<typename It, typename Compare = std::less<>>
bool next_combination(It lo, It mid, It hi, Compare comp = Compare{}) {
  using T = typename std::iterator_traits<It>::value_type;
  if (lo == mid || mid == hi) {
    return false;
  }
  It l = mid - 1, h = hi - 1;
  int len1 = 1, len2 = 1;
  while (l != lo && !comp(*l, *h)) {
    --l;
    ++len1;
  }
  if (l == lo && !comp(*l, *h)) {
    std::rotate(lo, mid, hi);
    return false;
  }
  for (; mid < h; ++len2) {
    if (!comp(*l, *--h)) {
      ++h;
      break;
    }
  }
  if (len1 == 1 || len2 == 1) {
    std::iter_swap(l, h);
  } else if (len1 == len2) {
    std::swap_ranges(l, mid, h);
  } else {
    std::iter_swap(l++, h++);
    int total = (--len1) + (--len2), gcd = total;
    for (int i = len1; i != 0;) {
      std::swap(gcd %= i, i);
    }
    int skip = total / gcd - 1;
    for (int i = 0; i < gcd; i++) {
      It curr = (i < len1) ? (l + i) : (h + (i - len1));
      int k = i;
      T prev = *curr;
      for (int j = 0; j < skip; j++) {
        k = (k + len1) % total;
        It next = (k < len1) ? (l + k) : (h + (k - len1));
        *curr = *next;
        curr = next;
      }
      *curr = prev;
    }
  }
  return true;
}

bool next_combination(int n, std::vector<int> &a) {
  int k = static_cast<int>(a.size());
  for (int i = k - 1; i >= 0; i--) {
    if (a[i] < n - k + i) {
      a[i]++;
      while (++i < k) {
        a[i] = a[i - 1] + 1;
      }
      return true;
    }
  }
  return false;
}

uint64_t next_combination_mask(uint64_t x) {
  if (x == 0) {
    return 0;
  }
  uint64_t s = x & -x, r = x + s;
  if (r == 0) {
    return 0;
  }
  return r | (((x ^ r) >> 2) / s);
}

int64_t n_choose_k(int64_t n, int64_t k) {
  if (k > n - k) {
    k = n - k;
  }
  int64_t res = 1;
  for (int i = 0; i < k; i++) {
    int64_t num = n - i, den = i + 1;
    int64_t g = std::gcd(num, den);
    num /= g;
    den /= g;
    res /= den;
    res *= num;  // Overflow warning.
  }
  return res;
}

std::vector<int> combination_by_rank(int n, int k, int64_t r) {
  std::vector<int> res(k);
  int count = n;
  for (int i = 0; i < k; i++) {
    int j = 1;
    for (;; j++) {
      int64_t am = n_choose_k(count - j, k - 1 - i);
      if (r < am) {
        break;
      }
      r -= am;
    }
    res[i] = (i > 0) ? (res[i - 1] + j) : (j - 1);
    count -= j;
  }
  return res;
}

int64_t rank_by_combination(int n, const std::vector<int> &a) {
  int k = static_cast<int>(a.size());
  int64_t res = 0;
  int prev = -1;
  for (int i = 0; i < k; i++) {
    for (int j = prev + 1; j < a[i]; j++) {
      res += n_choose_k(n - 1 - j, k - 1 - i);
    }
    prev = a[i];
  }
  return res;
}

bool next_combination_with_repeats(int n, std::vector<int> &a) {
  int k = static_cast<int>(a.size());
  for (int i = k - 1; i >= 0; i--) {
    if (a[i] < n - 1) {
      for (++a[i]; ++i < k;) {
        a[i] = a[i - 1];
      }
      return true;
    }
  }
  return false;
}

Example Usage

#include <cassert>
#include <iostream>
using namespace std;

template<typename It>
void print_range(It lo, It hi) {
  cout << "{";
  for (; lo != hi; ++lo) {
    cout << *lo << (lo == hi - 1 ? "" : ",");
  }
  cout << "} ";
}

int main() {
  {
    int k = 3;
    string s = "11234";
    cout << "\"" << s << "\" choose " << k << ":" << endl;
    int count = 0;
    do {
      cout << s.substr(0, k) << " ";
      count++;
    } while (next_combination(s.begin(), s.begin() + k, s.end()));
    assert(count == 7);  // Duplicate '1' values collapse equal output strings.
    cout << endl;
  }
  {  // Unordered combinations using masks.
    int n = 5, k = 3;
    string char_set = "abcde";  // Must be distinct.
    cout << "\n\"" << char_set << "\" choose " << k << " with masks:" << endl;
    uint64_t mask = (1ULL << k) - 1, limit = 1ULL << n;
    int count = 0;
    do {
      for (int i = 0; i < n; i++) {
        if ((mask >> i) & 1) {
          cout << char_set[i];
        }
      }
      cout << " ";
      count++;
      mask = next_combination_mask(mask);
    } while (mask < limit);
    assert(count == 10);
    assert(next_combination_mask(1ULL << 63) == 0);
    cout << endl;
  }
  {
    string s = "4321";
    assert(next_combination(s.begin(), s.begin() + 2, s.end(), greater<char>()));
    assert(s.substr(0, 2) == "42");
  }
  {  // Combinations of distinct integers in [0, n).
    int n = 5, k = 3;
    vector<int> a{0, 1, 2};
    cout << endl << n << " choose " << k << ":" << endl;
    int count = 0;
    do {
      print_range(a.begin(), a.end());
      assert(a == combination_by_rank(n, k, count));
      assert(rank_by_combination(n, a) == count);
      count++;
    } while (next_combination(n, a));
    assert(count == 10);
    assert(n_choose_k(66, 33) == 7219428434016265740LL);
    cout << endl;
  }
  {  // Combinations with repeats.
    int n = 3, k = 2;
    vector<int> a{0, 0};
    cout << endl << n << " multichoose " << k << ":" << endl;
    int count = 0;
    do {
      print_range(a.begin(), a.end());
      count++;
    } while (next_combination_with_repeats(n, a));
    assert(count == 6);
    cout << endl;
  }
  return 0;
}

Example Output

"11234" choose 3:
112 113 114 123 124 134 234

"abcde" choose 3 with masks:
abc abd acd bcd abe ace bce ade bde cde

5 choose 3:
{0,1,2} {0,1,3} {0,1,4} {0,2,3} {0,2,4} {0,3,4} {1,2,3} {1,2,4} {1,3,4} {2,3,4}

3 multichoose 2:
{0,0} {0,1} {0,2} {1,1} {1,2} {2,2}
Back to top