How to efficiently make a function call to each row of a 2D ndarray?

Question

I'm implementing a KNN classifier and need to quickly traverse the test set to calculate and store their predicted labels.

The way I use now is to use the list comprehension to get a list, and then turn it into a ndarray, similar to np.array([predict(point) for point in test_set]), but I think it takes time and space, because the for loop of Python is relatively slow and it needs to create another copy. Is there a more efficient way to get such an array? I know that numpy has apply_along_axis function, but it is said that it only implicitly uses the for loop, which may not improve the performance.

EDIT: I learned a possible way to save memory: match np.fromiter() function and generator, like np.fromiter((predict(point) for point in test_set), int, test_set.shape[0]), which avoids creating a list halfway. Unfortunately, in my program, it seems to run a little slower than the previous method.

Answer 1

the good old way:

def my_func(test_set):
  i = 0
  test_set_size = len(test_set)
  result = [None] * test_set_size
  while i < test_set_size:
    result[i] = predict(test_set[i])
    i = i + 1
  return np.array(result)