Use a list of values to select rows from a Pandas dataframe

Question:

Let’s say I have the following Pandas dataframe:

df = DataFrame({'A' : [5,6,3,4], 'B' : [1,2,3, 5]})
df

     A   B
0    5   1
1    6   2
2    3   3
3    4   5

I can subset based on a specific value:

x = df[df['A'] == 3]
x

     A   B
2    3   3

But how can I subset based on a list of values? – something like this:

list_of_values = [3,6]

y = df[df['A'] in list_of_values]

To get:

     A    B
1    6    2
2    3    3
Asked By: zach

||

Answers:

You can use the isin method:

In [1]: df = pd.DataFrame({'A': [5,6,3,4], 'B': [1,2,3,5]})

In [2]: df
Out[2]:
   A  B
0  5  1
1  6  2
2  3  3
3  4  5

In [3]: df[df['A'].isin([3, 6])]
Out[3]:
   A  B
1  6  2
2  3  3

And to get the opposite use ~:

In [4]: df[~df['A'].isin([3, 6])]
Out[4]:
   A  B
0  5  1
3  4  5
Answered By: Wouter Overmeire

You can use the method query:

df.query('A in [6, 3]')
# df.query('A == [6, 3]')

or

lst = [6, 3]
df.query('A in @lst')
# df.query('A == @lst')
Answered By: Mykola Zotko

Another method;

df.loc[df.apply(lambda x: x.A in [3,6], axis=1)]

Unlike the isin method, this is particularly useful in determining if the list contains a function of the column A. For example, f(A) = 2*A - 5 as the function;

df.loc[df.apply(lambda x: 2*x.A-5 in [3,6], axis=1)]

It should be noted that this approach is slower than the isin method.

Answered By: Achintha Ihalage

You can store your values in a list as:

lis = [3,6]

then

df1 = df[df['A'].isin(lis)]

Answered By: svp

list_of_values doesn’t have to be a list; it can be set, tuple, dictionary, numpy array, pandas Series, generator, range etc. and isin() and query() will still work.

Some common problems with selecting rows

1. list_of_values is a range

If you need to filter within a range, you can use between() method or query().

list_of_values = [3, 4, 5, 6] # a range of values

df[df['A'].between(3, 6)]  # or
df.query('3<=A<=6')

2. Return df in the order of list_of_values

In the OP, the values in list_of_values don’t appear in that order in df. If you want df to return in the order they appear in list_of_values, i.e. "sort" by list_of_values, use loc.

list_of_values = [3, 6]
df.set_index('A').loc[list_of_values].reset_index()

If you want to retain the old index, you can use the following.

list_of_values = [3, 6, 3]
df.reset_index().set_index('A').loc[list_of_values].reset_index().set_index('index').rename_axis(None)

3. Don’t use apply

In general, isin() and query() are the best methods for this task; there’s no need for apply(). For example, for function f(A) = 2*A - 5 on column A, both isin() and query() work much more efficiently:

df[(2*df['A']-5).isin(list_of_values)]         # or
df[df['A'].mul(2).sub(5).isin(list_of_values)] # or
df.query("A.mul(2).sub(5) in @list_of_values")

4. Select rows not in list_of_values

To select rows not in list_of_values, negate isin()/in:

df[~df['A'].isin(list_of_values)]
df.query("A not in @list_of_values")  # df.query("A != @list_of_values")

5. Select rows where multiple columns are in list_of_values

If you want to filter using both (or multiple) columns, there’s any() and all() to reduce columns (axis=1) depending on the need.

  1. Select rows where at least one of A or B is in list_of_values:
    df[df[['A','B']].isin(list_of_values).any(1)]
    df.query("A in @list_of_values or B in @list_of_values")
    
  2. Select rows where both of A and B are in list_of_values:
    df[df[['A','B']].isin(list_of_values).all(1)] 
    df.query("A in @list_of_values and B in @list_of_values")
    

Bonus:

You can also call isin() inside query():

df.query("A.isin(@list_of_values).values")
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.