How to create an np.array from another np.array if the element of a third np.array of the same index-CodePudding

Suppose I have

a = np.array([0, 1, 3, 0])
b = np.array([[4,5],[7,8],[9,5],[8,2]])

How can I create a third array

c = [[4,5],[8,2]]

with a list comprehension?

CodePudding user response：

c = np.array([b[i] for i in range(len(b)) if a[i] == 0])

CodePudding user response：

Note that the best (and fastest) way of doing this is without list comprehensions, and using boolean indexing:

c = b[a==0]

If you want to use list comprehensions, besides the option presented by @politinsa, you can use zip, that is useful when you have to go through more than one iterable at the same time (in this case, a and b):

c = np.array([b_d for (b_d, a_d) in zip(b, a) if a_d == 0])

When you want to iterate through elements and get the index of each element, enumerate can be cleaner than a range(len(iterable)) (although not necessarily faster):

c = np.array([b_d for i,b_d in enumerate(b) if a[i] == 0])