如何使用C++范围来实现numpy.ndindex?

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

本文关键字:numpy ndindex 实现 何使用 C++ 范围      更新时间:2023-10-16

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

下面是 2D 数组的实现。

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}).我可以想到多种解决方案,但我不确定哪种解决方案可以静态解决此问题。

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

这可以通过这种方式完成

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

现场示例