src/xbt/utils/iter/subsets.hpp

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