How to use C++ ranges to implement numpy.ndindex?

Question

I'd like to implement equivalent of numpy.ndindex in C++. It is supposed to generate indices for a multi-dimensional array of specified dimensions.

Here is the implementation for a 2D array.

template <typename T>
inline auto NDIndex(T d0, T d1) {
  using namespace ranges;
  return views::cartesian_product(views::indices(d0), views::indices(d1));
}

// Usage
for (const auto[i1, i2] : NDIndex(5, 4)) {
  arr[i1][i2] = ...
}

I'd like to generalize it for an arbitrary number of dimensions without sacrificing the performance. I'm ok with using brace in the interface, e.g. NDIndex({5, 4}). I can think of multiple solutions but I'm not sure which would resolve this statically.

Answer 1

views::cartesian_product is already variadic, you just need to expand a pack into it.

template <typename... Ts>
inline auto NDIndex(Ts ... ds) {
  using namespace ranges;
  return views::cartesian_product(views::indices(ds)...);
}

// Usage
int main() {
    for (const auto[i1, i2] : NDIndex(5, 4)) {
    }
    for (const auto[i1, i2, i3] : NDIndex(5, 4, 7)) {
    }
}

Answer 2

This can be done this way

#include <range/v3/view/indices.hpp>
#include <range/v3/view/cartesian_product.hpp>


template <unsigned... Ind>
constexpr inline auto NDIndex() {
  using namespace ranges;
  return views::cartesian_product(views::indices(Ind)...);
}


int main() {

    for (const auto[i1, i2] : NDIndex<5, 4>()) {
    }


    for (const auto[i1, i2, i3] : NDIndex<5, 4, 7>()) {
    }
}

Live example