Skip to content

Commit

Permalink
Add product_vec iterator
Browse files Browse the repository at this point in the history
- 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
  • Loading branch information
Olivier Parcollet authored and Wentzell committed Dec 9, 2020
1 parent 1a13d95 commit 22f2bbd
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 0 deletions.
82 changes: 82 additions & 0 deletions c++/itertools/itertools.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,45 @@ namespace itertools {
decltype(auto) dereference() const { return tuple_map_impl(std::index_sequence_for<It...>{}); }
};

/********************* Product Iterator with homegenous type but varying number of arg ********************/

// same logic as before, but at runtime
template <typename It>
struct prod_iter_vec : iterator_facade<prod_iter_vec<It>, std::vector<typename std::iterator_traits<It>::value_type>> {

std::vector<It> its_begin, its_end;
std::vector<It> its = its_begin;

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)) {}

void increment() {
for (int N = 0; N < its.size() - 1; ++N) {
++its[N];
if (its[N] != its_end[N]) return;
its[N] = its_begin[N];
}
++its[its.size() - 1];
}

bool equal(prod_iter_vec const &other) const { return (its == other.its); }

template <typename U>
bool equal(sentinel_t<U> const &s) const {
return (s.it == its.back());
}

template <typename U>
bool operator==(sentinel_t<U> const &s) const {
return equal(s);
}

std::vector<typename It::value_type> dereference() const {
std::vector<typename It::value_type> r(its.size());
for (int i = 0; i < its.size(); ++i) r[i] = *its[i];
return r;
}
};

/********************* Stride Iterator ********************/

template <typename Iter>
Expand Down Expand Up @@ -331,6 +370,39 @@ namespace itertools {

// ---------------------------------------------

template <typename T>
struct multiplied_vec {
std::vector<T> tu; // T can be a ref.

using iterator = prod_iter_vec<decltype(std::begin(std::declval<T &>()))>;
using const_iterator = prod_iter_vec<decltype(std::cbegin(std::declval<T &>()))>;

multiplied_vec(std::vector<T> const &ranges) : tu{ranges} {}

iterator begin() noexcept {
std::vector<typename T::iterator> _b(tu.size()), _e(tu.size());
std::transform(tu.begin(), tu.end(), _b.begin(), [](auto &&x) { return std::begin(x); });
std::transform(tu.begin(), tu.end(), _e.begin(), [](auto &&x) { return std::end(x); });
return iterator{_b, _e};
}

const_iterator cbegin() noexcept {
std::vector<typename T::const_iterator> _b(tu.size()), _e(tu.size());
std::transform(tu.begin(), tu.end(), _b.begin(), [](auto &&x) { return std::cbegin(x); });
std::transform(tu.begin(), tu.end(), _e.begin(), [](auto &&x) { return std::cend(x); });
return const_iterator{_b, _e};
}

auto end() noexcept { return make_sentinel(std::end(tu.back())); }
auto cend() const noexcept { return make_sentinel(std::cend(tu.back())); }
auto end() const noexcept { return cend(); }
};

template <typename T>
multiplied_vec(T &&) -> multiplied_vec<std::decay_t<T>>;

// ---------------------------------------------

template <typename T>
struct sliced {
T x;
Expand Down Expand Up @@ -458,6 +530,16 @@ namespace itertools {
return {std::forward<T>(ranges)...};
}

/**
* Lazy-product of multiple ranges. Same as product, but with an uniform type,
* but a number of ranges known at run time.
*
*/
template <typename T>
details::multiplied_vec<T> product_vec(std::vector<T> const &ranges) {
return {ranges};
}

/**
* Lazy-slice a range.
* This function returns itself a slice of the initial range
Expand Down
21 changes: 21 additions & 0 deletions test/c++/itertools.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,27 @@ TEST(Itertools, Product) {
EXPECT_EQ(V4, std::vector<int>(4, 1 * 2 * 3 * 4));
}

TEST(Itertools, ProductVec) {

std::vector<int> V1{0, 1, 2};
std::vector<int> V2{0, 1, 2};
std::vector<std::vector<int>> W{V1, V2};

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}};
for (auto vec : product_vec(W)) {
res.push_back(vec);
std::cout << "[" << vec[0] << "," << vec[1] << "]\n";
}

EXPECT_EQ(res.size(), check.size());
for (int i : range(res.size())) {
EXPECT_EQ(res[i].size(), check[i].size());
for (int j : range(res[i].size())) { EXPECT_EQ(res[i][j], check[i][j]); }
}

// make a real check
}

TEST(Itertools, Slice) {

for (long N : range(1, 6)) {
Expand Down

0 comments on commit 22f2bbd

Please sign in to comment.