Is there a Numpy equivalent of C++ std::vector reserve(), push_back() and shrink_to_fit()?

Question

I would like to append elements to en empty Numpy array in-place. I know the maximum array size beforehand. I can't find a direct way to accomplish that, so here is my workaround:

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)

Is there a more elegant way to code it, without keeping track of current length in j, similar in simplicity to C++ version below?

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();

And yes, I read How to extend an array in-place in Numpy?, but it doesn't cover this specific case, i.e. array size limit known beforehand.

Answer 1

Using np.fromiter(iter, dtype, count) with generator:

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

If the parameter count is omitted, the array size will expand automatically, but some performance will still be lost.

However, you cannot improve performance by specifying the parameter count as the maximum length, because if the iterator is not long enough, it will raise ValueError: iterator too short.