combinations.hpp

SECTIONUtility INCLUDEnoya/combinations.hpp

Enumerate all increasing k-subsets of [0,n) in lexicographic order, invoking callback(span-like vector); stop early when it returns false.

按字典序枚举 [0,n) 的所有 k 元子集，并允许回调提前停止。

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_COMBINATIONS_HPP
#define NOYA_COMBINATIONS_HPP 1

/// @complexity Time: O(k) per emitted k-subset.
/// Space: O(k).

#include <cassert>
#include <vector>

namespace noya {

/// @brief Enumerate all increasing k-subsets of [0,n) in lexicographic order,
/// invoking callback(span-like vector); stop early when it returns false.
template <class Callback>
void enumerate_combinations(int n, int k, Callback callback) {
  assert(0 <= k && k <= n);
  std::vector<int> chosen(k);
  for (int index = 0; index < k; index++) {
    chosen[index] = index;
  }
  if (k == 0) {
    callback(chosen);
    return;
  }
  while (true) {
    if (!callback(chosen)) {
      return;
    }
    int index = k - 1;
    while (index >= 0 && chosen[index] == n - k + index) {
      index--;
    }
    if (index < 0) {
      return;
    }
    chosen[index]++;
    for (int next = index + 1; next < k; next++) {
      chosen[next] = chosen[next - 1] + 1;
    }
  }
}

} // namespace noya

#endif // NOYA_COMBINATIONS_HPP

#include <cassert>
#include <vector>

/// @complexity Time: O(k) per emitted k-subset.
/// Space: O(k).

namespace noya {

/// @brief Enumerate all increasing k-subsets of [0,n) in lexicographic order,
/// invoking callback(span-like vector); stop early when it returns false.
template <class Callback>
void enumerate_combinations(int n, int k, Callback callback) {
  assert(0 <= k && k <= n);
  std::vector<int> chosen(k);
  for (int index = 0; index < k; index++) {
    chosen[index] = index;
  }
  if (k == 0) {
    callback(chosen);
    return;
  }
  while (true) {
    if (!callback(chosen)) {
      return;
    }
    int index = k - 1;
    while (index >= 0 && chosen[index] == n - k + index) {
      index--;
    }
    if (index < 0) {
      return;
    }
    chosen[index]++;
    for (int next = index + 1; next < k; next++) {
      chosen[next] = chosen[next - 1] + 1;
    }
  }
}

} // namespace noya