Skip to content

Commit 22f2bbd

Browse files
parcolletWentzell
authored andcommitted
Add product_vec iterator
- product_vec is similar to product, but takes a vector of ranges (hence homogeneous but of variable size). - Same logic as product, with run time decision - Add test
1 parent 1a13d95 commit 22f2bbd

File tree

2 files changed

+103
-0
lines changed

2 files changed

+103
-0
lines changed

c++/itertools/itertools.hpp

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,45 @@ namespace itertools {
199199
decltype(auto) dereference() const { return tuple_map_impl(std::index_sequence_for<It...>{}); }
200200
};
201201

202+
/********************* Product Iterator with homegenous type but varying number of arg ********************/
203+
204+
// same logic as before, but at runtime
205+
template <typename It>
206+
struct prod_iter_vec : iterator_facade<prod_iter_vec<It>, std::vector<typename std::iterator_traits<It>::value_type>> {
207+
208+
std::vector<It> its_begin, its_end;
209+
std::vector<It> its = its_begin;
210+
211+
prod_iter_vec(std::vector<It> its_begin, std::vector<It> its_end) : its_begin(std::move(its_begin)), its_end(std::move(its_end)) {}
212+
213+
void increment() {
214+
for (int N = 0; N < its.size() - 1; ++N) {
215+
++its[N];
216+
if (its[N] != its_end[N]) return;
217+
its[N] = its_begin[N];
218+
}
219+
++its[its.size() - 1];
220+
}
221+
222+
bool equal(prod_iter_vec const &other) const { return (its == other.its); }
223+
224+
template <typename U>
225+
bool equal(sentinel_t<U> const &s) const {
226+
return (s.it == its.back());
227+
}
228+
229+
template <typename U>
230+
bool operator==(sentinel_t<U> const &s) const {
231+
return equal(s);
232+
}
233+
234+
std::vector<typename It::value_type> dereference() const {
235+
std::vector<typename It::value_type> r(its.size());
236+
for (int i = 0; i < its.size(); ++i) r[i] = *its[i];
237+
return r;
238+
}
239+
};
240+
202241
/********************* Stride Iterator ********************/
203242

204243
template <typename Iter>
@@ -331,6 +370,39 @@ namespace itertools {
331370

332371
// ---------------------------------------------
333372

373+
template <typename T>
374+
struct multiplied_vec {
375+
std::vector<T> tu; // T can be a ref.
376+
377+
using iterator = prod_iter_vec<decltype(std::begin(std::declval<T &>()))>;
378+
using const_iterator = prod_iter_vec<decltype(std::cbegin(std::declval<T &>()))>;
379+
380+
multiplied_vec(std::vector<T> const &ranges) : tu{ranges} {}
381+
382+
iterator begin() noexcept {
383+
std::vector<typename T::iterator> _b(tu.size()), _e(tu.size());
384+
std::transform(tu.begin(), tu.end(), _b.begin(), [](auto &&x) { return std::begin(x); });
385+
std::transform(tu.begin(), tu.end(), _e.begin(), [](auto &&x) { return std::end(x); });
386+
return iterator{_b, _e};
387+
}
388+
389+
const_iterator cbegin() noexcept {
390+
std::vector<typename T::const_iterator> _b(tu.size()), _e(tu.size());
391+
std::transform(tu.begin(), tu.end(), _b.begin(), [](auto &&x) { return std::cbegin(x); });
392+
std::transform(tu.begin(), tu.end(), _e.begin(), [](auto &&x) { return std::cend(x); });
393+
return const_iterator{_b, _e};
394+
}
395+
396+
auto end() noexcept { return make_sentinel(std::end(tu.back())); }
397+
auto cend() const noexcept { return make_sentinel(std::cend(tu.back())); }
398+
auto end() const noexcept { return cend(); }
399+
};
400+
401+
template <typename T>
402+
multiplied_vec(T &&) -> multiplied_vec<std::decay_t<T>>;
403+
404+
// ---------------------------------------------
405+
334406
template <typename T>
335407
struct sliced {
336408
T x;
@@ -458,6 +530,16 @@ namespace itertools {
458530
return {std::forward<T>(ranges)...};
459531
}
460532

533+
/**
534+
* Lazy-product of multiple ranges. Same as product, but with an uniform type,
535+
* but a number of ranges known at run time.
536+
*
537+
*/
538+
template <typename T>
539+
details::multiplied_vec<T> product_vec(std::vector<T> const &ranges) {
540+
return {ranges};
541+
}
542+
461543
/**
462544
* Lazy-slice a range.
463545
* This function returns itself a slice of the initial range

test/c++/itertools.cpp

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,27 @@ TEST(Itertools, Product) {
105105
EXPECT_EQ(V4, std::vector<int>(4, 1 * 2 * 3 * 4));
106106
}
107107

108+
TEST(Itertools, ProductVec) {
109+
110+
std::vector<int> V1{0, 1, 2};
111+
std::vector<int> V2{0, 1, 2};
112+
std::vector<std::vector<int>> W{V1, V2};
113+
114+
std::vector<std::vector<int>> res, check{{0, 0}, {1, 0}, {2, 0}, {0, 1}, {1, 1}, {2, 1}, {0, 2}, {1, 2}, {2, 2}};
115+
for (auto vec : product_vec(W)) {
116+
res.push_back(vec);
117+
std::cout << "[" << vec[0] << "," << vec[1] << "]\n";
118+
}
119+
120+
EXPECT_EQ(res.size(), check.size());
121+
for (int i : range(res.size())) {
122+
EXPECT_EQ(res[i].size(), check[i].size());
123+
for (int j : range(res[i].size())) { EXPECT_EQ(res[i][j], check[i][j]); }
124+
}
125+
126+
// make a real check
127+
}
128+
108129
TEST(Itertools, Slice) {
109130

110131
for (long N : range(1, 6)) {

0 commit comments

Comments
 (0)