从pandas中的元组列表中选择行

时间:2014-06-05 15:52:21

标签: python pandas

我有DataFrame这样:

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

我想选择带有元组列表的行,如下所示:

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

我可以用我想要的东西:

# ( (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', ..])

这有效,但很尴尬和缓慢。 pandas中推荐的方式是什么?

1 个答案:

答案 0 :(得分:4)

只需将感兴趣的列设置为索引,排序并使用.loc

创建一个框架作为示例

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

使用元组选择

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