src/xbt/utils/iter/subsets.hpp

   1 /* Copyright (c) 2004-2023 The SimGrid Team. All rights reserved.           */
   2
   3 /* This program is free software; you can redistribute it and/or modify it
   4  * under the terms of the license (GNU LGPL) which comes with this package. */
   5
   6 #ifndef XBT_UTILS_ITER_SUBSETS_HPP
   7 #define XBT_UTILS_ITER_SUBSETS_HPP
   8
   9 #include <boost/iterator/iterator_facade.hpp>
  10 #include <functional>
  11 #include <numeric>
  12 #include <optional>
  13 #include <vector>
  14
  15 namespace simgrid::xbt {
  16
  17 /**
  18  * @brief A higher-order forward-iterator which traverses all possible subsets
  19  * of a given fixed size `k` of an iterable sequence
  20  *
  21  * @class Iterator: The iterator over which this higher-order iterator
  22  * generates elements.
  23  */
  24 template <class Iterator>
  25 struct subsets_iterator : public boost::iterator_facade<subsets_iterator<Iterator>, const std::vector<Iterator>,
  26                                                         boost::forward_traversal_tag> {
  27   subsets_iterator();
  28   explicit subsets_iterator(unsigned k);
  29   explicit subsets_iterator(unsigned k, Iterator begin, Iterator end = Iterator());
  30
  31 private:
  32   unsigned k; // The size of the subsets generated
  33   std::optional<Iterator> end = std::nullopt;
  34   std::vector<Iterator> current_subset;
  35   std::vector<unsigned> P; // Increment counts to determine which combinations need to be traversed
  36
  37   // boost::iterator_facade<...> interface to implement
  38   void increment();
  39   bool equal(const subsets_iterator<Iterator>& other) const;
  40   const std::vector<Iterator>& dereference() const;
  41
  42   // Allows boost::iterator_facade<...> to function properly
  43   friend class boost::iterator_core_access;
  44 };
  45
  46 template <typename Iterator> subsets_iterator<Iterator>::subsets_iterator() : subsets_iterator<Iterator>(0) {}
  47
  48 template <typename Iterator>
  49 subsets_iterator<Iterator>::subsets_iterator(unsigned k)
  50     : k(k), current_subset(std::vector<Iterator>(k)), P(std::vector<unsigned>(k))
  51 {
  52   std::iota(P.begin(), P.end(), k);
  53 }
  54
  55 template <typename Iterator>
  56 subsets_iterator<Iterator>::subsets_iterator(unsigned k, Iterator begin, Iterator end)
  57     : k(k), end(std::optional<Iterator>{end}), current_subset(std::vector<Iterator>(k)), P(std::vector<unsigned>(k))
  58 {
  59   for (unsigned i = 0; i < k; i++) {
  60     // Less than `k` elements to choose
  61     if (begin == end) {
  62       // We want to initialize the object then to be equivalent
  63       // to the end iterator so that there are no items to iterate
  64       // over
  65       this->end = std::nullopt;
  66       std::iota(P.begin(), P.end(), k);
  67       return;
  68     }
  69     current_subset[i] = begin++;
  70   }
  71   std::iota(P.begin(), P.end(), 0);
  72 }
  73
  74 template <typename Iterator> bool subsets_iterator<Iterator>::equal(const subsets_iterator<Iterator>& other) const
  75 {
  76   if (this->end == std::nullopt and other.end == std::nullopt) {
  77     return true;
  78   }
  79   if (this->k != other.k) {
  80     return false;
  81   }
  82   if (this->k == 0) { // this->k == other.k == 0
  83     return true;
  84   }
  85   return this->end != std::nullopt and other.end != std::nullopt and this->P[0] == other.P[0];
  86 }
  87
  88 template <typename Iterator> const std::vector<Iterator>& subsets_iterator<Iterator>::dereference() const
  89 {
  90   return this->current_subset;
  91 }
  92
  93 template <typename Iterator> void subsets_iterator<Iterator>::increment()
  94 {
  95   if (end == std::nullopt || k == 0) {
  96     return;
  97   }
  98
  99   // Move the last element over each time
 100   ++current_subset[k - 1];
 101   ++P[k - 1];
 102
 103   const auto end                  = this->end.value();
 104   const bool shift_other_elements = current_subset[k - 1] == end;
 105
 106   if (shift_other_elements) {
 107     if (k == 1) {
 108       // We're done in the case that k = 1; here, we've iterated
 109       // through the list once which is sufficient
 110       this->end = std::nullopt;
 111       return;
 112     }
 113
 114     // At this point, k >= 2
 115
 116     // The number of elements is now equal to the "index"
 117     // of the last element (it is at the end, which means we added
 118     // for the last time)
 119     const unsigned n = P[k - 1];
 120
 121     unsigned l = 0;
 122     for (unsigned j = k - 2; j > 0; j--) {
 123       if (P[j] != (n - (k - j))) {
 124         l = j;
 125         break;
 126       }
 127     }
 128
 129     ++P[l];
 130     ++current_subset[l];
 131
 132     // At this point, this means we've sucessfully iterated through
 133     // all subsets, so we're done
 134     if (l == 0 and P[0] > (n - k)) {
 135       this->end = std::nullopt;
 136       return;
 137     }
 138
 139     // Otherwise, everyone moves over
 140
 141     auto iter_at_l = current_subset[l];
 142     for (auto i = l + 1; i <= (k - 1); i++) {
 143       P[i]              = P[l] + (i - l);
 144       current_subset[i] = ++iter_at_l;
 145     }
 146   }
 147 }
 148
 149 } // namespace simgrid::xbt
 150
 151 #endif