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 && 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 && other.end != std::nullopt && 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 k == 0, there's nothing to do
  96   // If end == std::nullopt, we've finished
  97   // iterating over all subsets of size `k`
  98   if (end == std::nullopt || k == 0) {
  99     return;
 100   }
 101
 102   // Move the last element over each time
 103   ++current_subset[k - 1];
 104   ++P[k - 1];
 105
 106   const auto end                  = this->end.value();
 107   const bool shift_other_elements = current_subset[k - 1] == end;
 108
 109   if (shift_other_elements) {
 110     if (k == 1) {
 111       // We're done in the case that k = 1; here, we've iterated
 112       // through the list once, which is all that is needed
 113       this->end = std::nullopt;
 114       return;
 115     }
 116
 117     // At this point, k >= 2
 118
 119     // The number of elements is now equal to the "index"
 120     // of the last element (it is at the end, which means we added
 121     // for the last time)
 122     const unsigned n = P[k - 1];
 123
 124     // We're looking to determine
 125     //
 126     // argmax_{0 <= j <= k - 2}(P[j] != (n - (k - j)))
 127     //
 128     // If P[j] == (n - (k - j)) for some `j`, that means
 129     // the `j`th element of the current subset has moved
 130     // "as far as it can move" to the right; in other words,
 131     // this is our signal that some element before the `j`th
 132     // has to move over
 133     //
 134     // std::max_element() would work here too, but it seems
 135     // overkill to create a vector full of numbers when a simple
 136     // range-based for-loop can do the trick
 137     unsigned l = 0;
 138     for (unsigned j = k - 2; j > 0; j--) {
 139       if (P[j] != (n - (k - j))) {
 140         l = j;
 141         break;
 142       }
 143     }
 144
 145     ++P[l];
 146     ++current_subset[l];
 147
 148     // Plugging in `j = 0` to the above formula yields
 149     // `n - k`, which is the furthest point that the first (i.e. `0th`)
 150     // element can be located. Thus, if `P[0] > (n - k)`, this means
 151     // we've sucessfully iterated through all subsets so we're done
 152     if (P[0] > (n - k)) {
 153       this->end = std::nullopt;
 154       return;
 155     }
 156
 157     // Otherwise, all elements past element `l` are reset
 158     // to follow one after another immediately after `l`
 159     auto iter_at_l = current_subset[l];
 160     for (auto i = l + 1; i <= (k - 1); i++) {
 161       P[i]              = P[l] + (i - l);
 162       current_subset[i] = ++iter_at_l;
 163     }
 164   }
 165 }
 166
 167 } // namespace simgrid::xbt
 168
 169 #endif