是否有 C++ std::vector Reserve()、push_back() 和 Shrink_to_fit() 的 Numpy 等效项？

Question

我想将元素添加到空 Numpy 数组中。我事先知道最大数组大小。我找不到直接的方法来实现这一点，所以这是我的解决方法：

N = 1000
a = np.empty([N], dtype=np.int32)
j = 0

for i in range(N):
  if f(i):
    a[j] = g(i)
    j += 1

a.resize(j)

是否有更优雅的编码方式，无需跟踪 中的当前长度j，类似于下面的 C++ 版本的简单性？

const int N = 1000;

vector<int> a;
a.reserve(N);

for (int i=0; i<N; i++) 
  if (f(i)) 
    a.push_back(g(i));

a.shrink_to_fit();

是的，我读过如何在 Numpy 中就地扩展数组？，但它不涵盖这种特定情况，即事先已知的数组大小限制。

Answer 1

np.fromiter(iter, dtype, count)与发电机一起使用：

a = np.fromiter((g(i) for i in range(N) if f(i)), np.int32)

如果省略该参数count，数组大小会自动扩展，但仍然会损失一些性能。

但是，您无法通过将参数指定count为最大长度来提高性能，因为如果迭代器不够长，它将引发ValueError: iterator too short。