有没有 Numpy 相当于 C++ std::vector 的 reserve()、push_back() 和 shrink_to_fit() 函数？

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中就地扩展数组？，但它并没有涵盖这种特定情况，即事先已知数组大小限制。

- Paul Jurczak

---

为什么要建立一个数组并进行调整大小？你可以直接使用列表推导式 a = np.array([g(i) for i in range(N) if f(i)])。 - Ch3steR

Python不是带有风格的C++，当你进行低级别的操作时，它会变得很痛苦，如果可能的话。 - Passer By

resize方法是唯一可以原地改变numpy数组大小的操作。 - hpaulj

在提供任何建议之前，我需要更多关于使用情况的信息。N 实际上会有多大？你可以尝试两个版本：一个是像你现在所做的那样在最后调整大小，执行一次复制，另一个版本则先循环一遍来确定实际大小。测量性能并相应地选择。 - codie

@PaulJurczak 对我来说，这不是关于我是否喜欢某件事。而是关于适当地解决技术问题。如果通过添加另一个循环来解决用例问题，那就这样做吧。但你最了解你的用例。 - codie

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错误。

- Mechanic Pig