Search code examples
pythonnumpylist-comprehension

How to create an np.array from another np.array if the element of a third np.array of the same index is zero?

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?

Solution

  • 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])