python - 使用计数将条件排除应用于 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 方法冲突。