过滤掉没有足够数量的满足条件的行的组

Question

我有以下熊猫数据框。

import pandas as pd

# Initialize dataframe
df1 = pd.DataFrame(columns=['bar', 'foo'])
df1['bar'] = ['001', '001', '001', '001', '002', '002', '003', '003', '003']
df1['foo'] = [-4, -3, 2, 3, -3, -2, 0, 1, 2]
>>> print df1
   bar  foo
0  001   -4
1  001   -3
2  001    2
3  001    3
4  002   -3
5  002   -2
6  003    0
7  003    1
8  003    2

考虑以下阈值和参数。

# Provide threshold and number of entries above and below threshold
threshold = 0
n_below = 2
n_above = 2

我想创建一个数据框来过滤掉 的某些值bar。我想过滤掉的内容bar如下：如果它没有至少小于的n_below值和大于的值。foothresholdn_abovefoothreshold

对于上面的例子：

• 该组bar = 001不会被过滤掉，因为bar = 001至少存在小于的n_below = 2条目和至少大于的条目。foothreshold = 0n_above = 2foothreshold = 0
• 该组bar = 002将被过滤掉，因为bar = 002至少不存在大于 的n_above = 2条目。foothreshold = 0
• 该组bar = 003将被过滤掉，因为bar = 003至少不存在小于 的n_below = 2条目。foothreshold = 0

所需的输出如下：

# Desired output
   bar  foo
0  001   -4
1  001   -3
2  001    2
3  001    3

我相信这可以通过 GroupBy 和 来完成.count()，但是我一直无法找到有效的解决方案。我认识到编写一个分两步执行此操作的解决方案可能会更清晰：1）首先过滤以满足条件n_below；2）然后过滤满足条件n_above。

Answer 1

您可以使用groupby和filter方法。

threshold = 0
n_below = 2
n_above = 2
def filter_function(g):
    '''Called by filter, g is the grouped dataframe'''
    l = g['foo']
    return  (sum([x < threshold for x in l]) >= n_below 
            and sum([x > threshold for x in l]) >= n_above)

df.groupby('bar').filter(filter_function)

# gives
    bar foo
0   1   -4
1   1   -3
2   1   2
3   1   3

参见Pandas：过滤