Select rows of a dataframe according to the elements of an array in one column

user_6372859 0
Viewed 1628

I have a dataframe in which one of the columns has one-dimensional arrays as values. As a toy example:

easy={"a":[1,2,3,4,5,6,7,8,9],"b":[9,8,7,6,5,4,3,2,1], "c":[[0.9,0.3,0.1],[0.8,0.7,0.2],[0.7,0.6,0.3],
                                                            [0.6,0.2,0.4],[0.5,0.9,0.5],[0.4,0.8,0.6],
                                                            [0.3,0.5,0.7],[0.2,0.1,0.8],[0.1,0.4,0.9]]}

easy_df=pd.DataFrame(easy) 

    a   b          c
0   1   9   [0.9, 0.3, 0.1]
1   2   8   [0.8, 0.7, 0.2]
2   3   7   [0.7, 0.6, 0.3]
3   4   6   [0.6, 0.2, 0.4]
4   5   5   [0.5, 0.9, 0.5]
5   6   4   [0.4, 0.8, 0.6]
6   7   3   [0.3, 0.5, 0.7]
7   8   2   [0.2, 0.1, 0.8]
8   9   1   [0.1, 0.4, 0.9]

The arrays in column "c" could have in principle many elements. I would like to make another dataframe out of this one, but only with those rows such as any of the elements of the arrays in column "c" is higher than a certain threshold, for example: 0.75. The new dataframe should look like:

    a   b          c
0   1   9   [0.9, 0.3, 0.1]
1   2   8   [0.8, 0.7, 0.2]
4   5   5   [0.5, 0.9, 0.5]
5   6   4   [0.4, 0.8, 0.6]
7   8   2   [0.2, 0.1, 0.8]
8   9   1   [0.1, 0.4, 0.9]

because these are the only rows of the initial dataframe that the arrays in "c" contain at least one element higher than 0.75.

I tried something like this:

easy_df[np.any(easy_df["c"])>0.75]

which is obviously wrong, as I get the error message 

TypeError: '>' not supported between instances of 'list' and 'float'

PS: The columns "a" and "b" have nothing to do with the problem, as they could have arbitrary elements.

pandas dataframe arrays
3 Answers

You can Convert the series of list into a dataframe and check if any of the column is greater than the threshold:

thresh = 0.75
easy_df[pd.DataFrame(easy_df['c'].tolist()).gt(thresh).any(1)]

Or using explode with aggregation of any on level=0:

df[easy_df['c'].explode().gt(thresh).any(level=0)]

   a  b                c
0  1  9  [0.9, 0.3, 0.1]
1  2  8  [0.8, 0.7, 0.2]
4  5  5  [0.5, 0.9, 0.5]
5  6  4  [0.4, 0.8, 0.6]
7  8  2  [0.2, 0.1, 0.8]
8  9  1  [0.1, 0.4, 0.9]

EDIT: Per comment , you can add an idxmax on axis=1:

thresh = 0.75
ref = pd.DataFrame(easy_df['c'].tolist())
out = easy_df[ref.gt(thresh).any(1)].assign(d=ref.idxmax(1))
print(out)

   a  b                c  d
0  1  9  [0.9, 0.3, 0.1]  0
1  2  8  [0.8, 0.7, 0.2]  0
4  5  5  [0.5, 0.9, 0.5]  1
5  6  4  [0.4, 0.8, 0.6]  1
7  8  2  [0.2, 0.1, 0.8]  2
8  9  1  [0.1, 0.4, 0.9]  2
user_98406370

IIUC,

easy_df.loc[(np.stack(easy_df['c']) > 0.75).any(axis=1)]

   a  b                c
0  1  9  [0.9, 0.3, 0.1]
1  2  8  [0.8, 0.7, 0.2]
4  5  5  [0.5, 0.9, 0.5]
5  6  4  [0.4, 0.8, 0.6]
7  8  2  [0.2, 0.1, 0.8]
8  9  1  [0.1, 0.4, 0.9]

Time for sample DataFrame

%%timeit
easy_df.loc[(np.stack(easy_df['c']) > 0.75).any(axis=1)]
483 µs ± 13 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
user_118842370

Get the max value from each list in c, compare to your threshold 0.75

easy_df[easy_df['c'].apply(np.max) > 0.75]
user_32207690
Related