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

Question

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

这是二维数组的实现。

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

Answer 1

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