Fraz Fraz - 11 days ago 5
Python Question

Better way to vectorize using numpy

I am wondering if there is a better way to vectorize the following:

array = []
for ele in long_list:
if condition:
array.append(1)
vector = np.array(array)


Now this is essentially
np.ones(len(long_list))
except for that
if
condition.

Is there a better (vectorized) way of achieving the above?

Answer
vector = np.fromiter((1 for ele in long_list if condition), dtype=int)

This should be faster than either:

vector = np.ones(len([x for x in long_list if conditon]))

or

vector = np.array( [1 for ele in long_list if condition(ele)] )

because it avoids building the filtered list, due to the fact that (1 for ele in long_list if condition) returns a generator not a list.

EDIT

I'm wondering if this isn't also a viable option:

vector = np.array( (1 for ele in long_list if condition(ele)) )