Select rows from a list of tuples in pandas

Question

I have a DataFrame like this:

In [140]: df.head()
Out[140]: 
   user  brand  action  nthday
0   767   3961       0      51
1   767   3961       2      51
2   767   3961       2      51
3   767   3961       0      51
4   767   3961       0      51

I want to select rows with a list of tuple like this:

mylist = [(767, 3961), (768, 4201),...]

I can do with what I want with:

# ( (767, 3961, 0, 51), ... )
intermediate_ = ( tuple(r) for (i,r) in df.iterrows() if tuple(r)[:2] in set(mylist))

# reconstruct a DataFrame
subdf = DataFrame(intermediate_, columns = ['user', 'brand', 'action', ..])

This works but is awkward and slow. What is the recommended way in pandas?

Answer 1

Just set the columns of interest as an index, sort, and use .loc

Create a frame as an example

In [8]: df = DataFrame(np.random.randn(12,1),index=pd.MultiIndex.from_product([list(range(3)),list(range(4))],names=['foo','bar']))

In [10]: df.reset_index()
Out[10]: 
    foo  bar         0
0     0    0 -0.225873
1     0    1 -0.275865
2     0    2 -1.324766
3     0    3 -0.607122
4     1    0 -1.465992
5     1    1 -1.582276
6     1    2 -0.718533
7     1    3 -1.904252
8     2    0  0.588496
9     2    1 -1.057599
10    2    2  0.388754
11    2    3 -0.940285

In [11]: x = df.reset_index()

In [12]: df2 = x.set_index(['foo','bar']).sort_index()

In [13]: df2
Out[13]: 
                0
foo bar          
0   0   -0.225873
    1   -0.275865
    2   -1.324766
    3   -0.607122
1   0   -1.465992
    1   -1.582276
    2   -0.718533
    3   -1.904252
2   0    0.588496
    1   -1.057599
    2    0.388754
    3   -0.940285

Select with tuples

In [14]: df2.loc[[(0,2),(2,0),(2,3)]]
Out[14]: 
                0
foo bar          
0   2   -1.324766
2   0    0.588496
    3   -0.940285