Filter dictionary whose values are arrays

user_11970084 0
Viewed 78

I have data which looks like this:

features_dict = {
    'feat1': np.array([[0,1],[2,3],[4,5]]), 
    'feat2': np.array([[6,7],[8,9],[10,11]]),
    'feat3': np.array([1, 0, 0]),
    'feat4': np.array([[1],[2],[1]])
}

I want to filter the values of above dictionary based on the first dimension index where feat3 values are 0. Hence, the output I'm looking for is:

features_dict = {
    'feat1': np.array([[2,3],[4,5]]), 
    'feat2': np.array([[8,9],[10,11]]),
    'feat3': np.array([0, 0]),
    'feat4': np.array([[2],[1]])
}

Notice that I want to have only the 2nd and 3rd elements of each dict value because that's where feat3 values are 0.

Initially, I was thinking of converting the dict to pandas and filter the rows using .loc but it turned out that pandas can't accept arrays.

Can anyone please help? Thanks

numpy python
1 Answers
import numpy as np

features_dict = {
    'feat1': np.array([[0,1],[2,3],[4,5]]), 
    'feat2': np.array([[6,7],[8,9],[10,11]]),
    'feat3': np.array([1, 0, 0]),
    'feat4': np.array([[1],[2],[1]])
}

ind = features_dict['feat3'] == 0
features_dict = {k: v[ind] for k,v in features_dict.items()}

After filtering:

{
'feat1': array([[2, 3],[4, 5]]),
 'feat2': array([[ 8,  9],[10, 11]]),
 'feat3': array([0, 0]),
 'feat4': array([[2],[1]])
}
user_32796030
Related