Subset numpy array based on another array

Viewed 5090

I have 2 numpy arrays:

arr_a = array(['1m_nd', '2m_nd', '1m_4wk'],
      dtype='<U15')

arr_b = array([0, 1, 1])

I want to select elements from arr_a based on arr_b. I am doing this: arr_a[arr_b], but I get this as result:

array(['1m_nd', '2m_nd', '2m_nd'],
      dtype='<U15')

instead of:

array(['2m_nd', '1m_4wk'],
      dtype='<U15')

How do i fix this?

2 Answers

You need to pass it a boolean array, for example:

>>> arr_a[arr_b>0]

array(['2m_nd', '1m_4wk'], 
      dtype='<U15')

Given arr_a and arr_b, Running the following will give the boolean array for each of the elements in arr_b whose value is 1 => True and 0 => False . Correspondingly the boolean values are checked with the index value in arr_a. Here is the line of code you'd need.

>>> arr_a[arr_b == 1]
array([u'2m_nd', u'1m_4wk'],
      dtype='<U15')
Related