How to filter a NumPy array using value of another NumPy Array

Question

I have one NumPy array like this:

type = array(["A", "B", "A", "B", "B", "A", "A"]) 

I want to use this array to filter another array like this:

data = array([5,4,5], [0,4,3], [2,1,6], [6,7,9], [0,1,4], [7,9,8], [1,4,9]) 

I want to be able to filter the data by type, such that:

typeA = array([5,4,5],[2,1,6],[7,9,8],[1,4,9]) typeB = array([0,4,3],[6,7,9],[0,1,4]) 

Any idea how should I implement it? Since the first array is not boolean, I'm having a lot of trouble to figure out the solution

Answer 1

Let's define your arrays:

>>> from numpy import array >>> data = array(([5,4,5], [0,4,3], [2,1,6], [6,7,9], [0,1,4], [7,9,8], [1,4,9])) >>> type_ = array(["A", "B", "A", "B", "B", "A", "A"]) 

Now, let's find your typeA:

>>> A = data[type_=="A"] >>> A array([[5, 4, 5], [2, 1, 6], [7, 9, 8], [1, 4, 9]]) 

And your typeB:

>>> B = data[type_=="B"] >>> B array([[0, 4, 3], [6, 7, 9], [0, 1, 4]]) 

Here, type_=="A" or type_=="A" are boolean arrays. We use them to select the elements from data that we want.

For more on indexing with boolean or “mask” arrays, see the NumPy docs.

Lastly, type is a python builtin and it is best practices not to overwrite it. So, we used type_ here instead.

Answer 2

import numpy as np type_ = np.array(['A','B','A','B','B','A','A']) data = np.array([[5,4,5], [0,4,3], [2,1,6], [6,7,9], [0,1,4], [7,9,8], [1,4,9]]) from collections import defaultdict type_splitted = defaultdict(list) for type_i, data_i in zip(type_, data): type_splitted[type_i].append(data_i) typeA = np.asarray(type_splitted['A']) typeB = np.asarray(type_splitted['B'])