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.

Caleth · Accepted Answer · 2020-01-20T15:09:56.213

4

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

edited Jan 20 '20 at 15:09

answered Jan 20 '20 at 11:25

Caleth

52,200
2
44
75

@bartop oops, stole your examples without changing them :p – Caleth Jan 20 '20 at 11:30
btw, still does not compile :P – bartop Jan 20 '20 at 13:50

bartop · Answer 2 · 2020-01-20T11:30:21.047

3

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

edited Jan 20 '20 at 11:30

answered Jan 20 '20 at 11:21

bartop

9,971
1
23
54

I didn't mention that, but the limits can not be constexpr. I don't know them in advance. – pkubik Jan 20 '20 at 15:11

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

2 Answers2