使用计数将条件排除应用于 Pandas DataFrame

Question

我在 pandas 中有以下 DataFrame：

import pandas as pd
example_data = [{'ticker': 'aapl', 'loc': 'us'}, {'ticker': 'mstf', 'loc': 'us'}, {'ticker': 'baba', 'loc': 'china'}, {'ticker': 'ibm', 'loc': 'us'}, {'ticker': 'db', 'loc': 'germany'}]
df = pd.DataFrame(example_data)
print df

loc ticker
0       us   aapl
1       us   mstf
2    china   baba
3       us    ibm
4  germany     db

我想创建一个新的 DataFrame，这样每一行都是从原始 df 创建的，但 loc 计数大于 2 的行被排除在外。也就是说，新的 df 是通过循环旧的 df，计算之前出现的 loc 行数，并根据此计数包括/排除行来创建的。

以下代码给出了所需的输出。

country_counts = {}
output = []
for row in df.values:
    if row[0] not in country_counts:
        country_counts[row[0]] = 1
    else:
        country_counts[row[0]] +=1
    if country_counts[row[0]] <= 2:
        output.append({'loc': row[0], 'ticker': row[1]})
new_df = pd.DataFrame(output)   
print new_df

loc ticker
0       us   aapl
1       us   mstf
2    china   baba
3  germany     db

输出不包括原始 df 中的第 4 行，因为它的 loc 计数大于 2（即 3）。

是否有更好的方法来执行此类操作？非常感谢任何帮助。

Answer 1

groupby 和 .head:

In [90]: df.groupby('loc').head(2)
Out[90]: 
       loc ticker
0       us   aapl
1       us   mstf
2    china   baba
4  germany     db

此外，请注意您的列名，因为 loc 与 .loc 方法冲突。

使用计数将条件排除应用于 Pandas DataFrame

Applying Conditional Exclusions to Pandas DataFrame using Counts

python

python-2.7

pandas