filter dataframe by list of params

Viewed 82

Got a DataFrame like this:

pd.DataFrame({"type":["type1","type1","type1","type2","type2","type2","type1/type2","type1/type2","type1/type2"],
                    "id": [1,2,3,1,2,3,1,2,3],
                    "values":[1,np.nan,7,np.nan,5,np.nan,1,1,1]})

:

type id values
type1 1 1
type1 2 NaN
type1 3 7
type2 1 NaN
type2 2 5
type2 3 NaN
type1,type2 1 1
type1,type2 2 1
type1,type2 3 1

And I have a "guide" like this (types with ids they corresponds to):

type1 = [1,3]
type2 = [2,3]
type1,type2 = [1,2,3]

I'd like to filter rows based on this lists (you can see that there're two cases possible: type1-id2 = NaN need to be filtered, but type2-id3=NaN - is a valid value)

How could i perform this operation in a best way? I have much bigger table and more guide lists in fact.

3 Answers

I suggest create dictionary and then selecting with np.logical_or.reduce with loop by values in dictionary with Series.eq for == and Series.isin for check membership, last pass to boolean indexing:

d = {}
d['type1'] = [1,3]
d['type2'] = [2,3]
d['type1/type2'] = [1,2,3]
print (d)
{'type1': [1, 3], 'type2': [2, 3], 'type1/type2': [1, 2, 3]}

m = np.logical_or.reduce([df['type'].eq(k) & df['id'].isin(v) for k, v in d.items()])

df = df[m]
print (df)
          type  id  values
0        type1   1     1.0
2        type1   3     7.0
4        type2   2     5.0
5        type2   3     NaN
6  type1/type2   1     1.0
7  type1/type2   2     1.0
8  type1/type2   3     1.0

Another aproach is create new DataFrame and merge by all intersected columns by original, here type and id columns:

d = {}
d['type1'] = [1,3]
d['type2'] = [2,3]
d['type1/type2'] = [1,2,3]
print (d)

df1 = pd.DataFrame([(k, x) for k, v in d.items() for x in v], columns=['type','id'])
print (df1)
          type  id
0        type1   1
1        type1   3
2        type2   2
3        type2   3
4  type1/type2   1
5  type1/type2   2
6  type1/type2   3

df = df.merge(df1)
print (df)
          type  id  values
0        type1   1     1.0
1        type1   3     7.0
2        type2   2     5.0
3        type2   3     NaN
4  type1/type2   1     1.0
5  type1/type2   2     1.0
6  type1/type2   3     1.0

Create a map from type to the valid ids:

d = {'type1':[1, 3], 'type2':[2, 3], 'type1/type2': [1, 2, 3]}

Then create boolean map for the ids to keep:

keep_bool = [id in d[type] for type, id in zip(df['type'], df['id'])]

Then only keep the valid indices.

df = df[keep_bool]

Input:

        type  id  values
0        type1   1     1.0
1        type1   2     NaN
2        type1   3     7.0
3        type2   1     NaN
4        type2   2     5.0
5        type2   3     NaN
6  type1/type2   1     1.0
7  type1/type2   2     1.0
8  type1/type2   3     1.0

Output:

>>> df[keep_bool]
          type  id  values
0        type1   1     1.0
2        type1   3     7.0
4        type2   2     5.0
5        type2   3     NaN
6  type1/type2   1     1.0
7  type1/type2   2     1.0
8  type1/type2   3     1.0

You could use a combination of itertools, as well as indexing to filter the values you want:

from itertools import product, chain

d = {'type1':[1, 3], 'type2':[2, 3], 'type1/type2': [1, 2, 3]}

combinations = chain.from_iterable(product([key], value) 
                                   for key, value 
                                   in d.items()
                                   )

    df.set_index(['type', 'id']).loc[combinations].reset_index()
 
          type  id  values
0        type1   1     1.0
1        type1   3     7.0
2        type2   2     5.0
3        type2   3     NaN
4  type1/type2   1     1.0
5  type1/type2   2     1.0
6  type1/type2   3     1.0
Related