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 // template <class Iterable> struct LazyPowerSet {
  19 // };
  20
  21 // template <class Iterable> struct LazyKSubsets {
  22 // private:
  23 //   const Iterable& universe;
  24
  25 // public:
  26 //   LazyKSubsets(const Iterable& universe) : universe(universe) {}
  27 // };
  28
  29 template <class Iterator> struct subsets_iterator {
  30
  31   subsets_iterator(unsigned k);
  32   subsets_iterator(unsigned k, Iterator begin, Iterator end = Iterator());
  33
  34   subsets_iterator& operator++();
  35   auto operator->() const { return &current_subset; }
  36   auto& operator*() const { return current_subset; }
  37
  38   bool operator==(const subsets_iterator<Iterator>& other) const
  39   {
  40     if (this->k != other.k) {
  41       return false;
  42     }
  43     if (this->k == 0) {
  44       return true;
  45     }
  46     return this->P[0] == other.P[0];
  47   }
  48   bool operator!=(const subsets_iterator<Iterator>& other) const { return not(this->operator==(other)); }
  49
  50   using iterator_category = std::forward_iterator_tag;
  51   using difference_type   = int; // # of steps between
  52   using value_type        = std::vector<Iterator>;
  53   using pointer           = value_type*;
  54   using reference         = value_type&;
  55
  56 private:
  57   unsigned k;
  58   std::optional<Iterator> end = std::nullopt;
  59   std::vector<Iterator> current_subset;
  60   std::vector<unsigned> P;
  61 };
  62
  63 template <typename Iterator>
  64 subsets_iterator<Iterator>::subsets_iterator(unsigned k)
  65     : k(k), current_subset(std::vector<Iterator>(k)), P(std::vector<unsigned>(k))
  66 {
  67   std::iota(P.begin(), P.end(), k);
  68 }
  69
  70 template <typename Iterator>
  71 subsets_iterator<Iterator>::subsets_iterator(unsigned k, Iterator begin, Iterator end)
  72     : k(k), end(std::optional<Iterator>{end}), current_subset(std::vector<Iterator>(k)), P(std::vector<unsigned>(k))
  73 {
  74   for (unsigned i = 0; i < k; i++) {
  75     // Less than `k` elements to choose
  76     if (begin == end) {
  77       // We want to initialize the object then to be equivalent
  78       // to the end iterator so that there are no items to iterate
  79       // over
  80       this->end = std::nullopt;
  81       std::iota(P.begin(), P.end(), k);
  82       return;
  83     }
  84     current_subset[i] = begin++;
  85   }
  86   std::iota(P.begin(), P.end(), 0);
  87 }
  88
  89 template <typename Iterator> subsets_iterator<Iterator>& subsets_iterator<Iterator>::operator++()
  90 {
  91   if (end == std::nullopt || k == 0) {
  92     return *this;
  93   }
  94
  95   // Move the last element over each time
  96   ++current_subset[k - 1];
  97   ++P[k - 1];
  98
  99   const auto end                  = this->end.value();
 100   const bool shift_other_elements = current_subset[k - 1] == end;
 101
 102   if (shift_other_elements) {
 103     // The number of elements is now equal to the "index"
 104     // of the last element (it is at the end, which means we added
 105     // for the finalth time)
 106     const unsigned n = P[k - 1];
 107
 108     auto l = 0;
 109     for (int j = static_cast<int>(k - 2); j >= 0; j--) {
 110       if (P[j] != (n - (k - j))) {
 111         l = j;
 112         break;
 113       }
 114     }
 115
 116     P[l] += 1;
 117
 118     if (l == 0 and P[0] > (n - k)) {
 119       return *this;
 120     }
 121
 122     for (auto i = l + 1; i <= (k - 1); i++) {
 123       P[i] = P[l] + (i - l);
 124     }
 125   }
 126
 127   return *this;
 128 }
 129
 130 } // namespace simgrid::xbt
 131
 132 #endif