如何使用C++ ranges实现numpy.ndindex?

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

我想在 C++ 中实现 numpy.ndindex 的等价物。它应该为指定维度的多维数组生成索引。

下面是二维数组的实现。

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] = ...
}

我想在不牺牲性能的情况下将其概括为任意数量的维度。我可以在界面中使用大括号,例如NDIndex({5, 4})。我可以想到多种解决方案,但不确定哪个可以静态解决此问题。

可以这样做

#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

views::cartesian_product已经是可变的,你只需要扩展一个包进去。

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