pandas：具有条件公式和唯一值的 groupby答案

【问题标题】：pandas: groupby with conditional formula and unique valuespandas：具有条件公式和唯一值的 groupby
【发布时间】：2021-10-25 01:28:54
【问题描述】：

我有以下数据框

import pandas as pd
  
country = ['US', 'US', 'US', 'UK', 'UK', 'Canada', 'Canada', "Mexico"]
feature =  [2, 2, 2, 1, 1, 2, 2, 1]
ID = [1, 2, 1, 3, 4, 1, 2, 1]  

df = pd.DataFrame(list(zip(country,feature, ID)),
               columns =['country', 'feature', 'ID'])

这是

   country  feature ID
0   US        2     1
1   US        2     2
2   US        2     1
3   UK        1     3
4   UK        1     4
5   Canada    2     1
6   Canada    2     2
7   Mexico    1     1

鉴于上述数据框，我想在 country 上执行 groupby。我想使用以下规则为每个country 创建一个值count：如果feature ==2，count = number of unique entries under ID for that country。如果feature !=2，count = number of total entries for that country。因此，生成的新数据框将如下所示：

  country   count
0   US        2
1   UK        2
2   Canada    2
3   Mexico    1

【问题讨论】：

为什么是加拿大1？
如果某些国家/地区同时拥有feature == 2 和feature != 2 怎么办？

标签： python pandas dataframe conditional-statements pandas-groupby

【解决方案1】：

您可以在两个groupby 操作中完成此操作，这实际上比一个包含lambda 的groupby 来测试feature 要快。此外，为了解决某些国家/地区可能意外混合feature == 2 和feature != 2 的情况，我们将取最大值：

a = df.groupby(['country', 'ID']).size().groupby('country').size()
b = df.groupby('country').agg({'feature': max, 'ID': 'size'})
out = b['ID'].where(b['feature'] != 2, a).to_frame('count')

>>> out
         count
country       
Canada       2
Mexico       1
UK           2
US           2

速度比较

# 1. setup
n = 100_000
ci = np.random.randint(0, 80, n)
df = pd.DataFrame({
    'country': [f'c_{i}' for i in ci],
    'feature': [(i % 2) + 1 for i in ci],
    'ID': np.random.randint(0, 60, n),
})

# test
%timeit ours(df)
# 17.5 ms ± 153 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit with_lambda(df)
# 34.7 ms ± 49.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

【讨论】：

【解决方案2】：

只需在groupby内做一个简单的条件

out = df.groupby('country').\
           apply(lambda x :x['ID'].nunique() if x['feature'].eq(2).all() else len(x))
Out[649]: 
country
Canada    2
Mexico    1
UK        2
US        2
dtype: int64

【讨论】：

【解决方案3】：

假设每个国家/地区只有一个值feature，您可以使用.nunique() 获取feature==2 的唯一计数，并使用.size() 获取feature==1 的总条目数。然后，组合结果并按国家/地区排序：

df_eq = df.loc[df['feature'] == 2].groupby('country')['ID'].nunique().reset_index(name='count')
df_ne = df.loc[df['feature'] != 2].groupby('country')['ID'].size().reset_index(name='count')

df_out = df_eq.append(df_ne).sort_values('country').reset_index(drop=True)

结果：

print(df_out)

  country  count
0  Canada      2
1  Mexico      1
2      UK      2
3      US      2

【讨论】：

【解决方案4】：

Canada 的计数在预期输出中应为 2。

您可以按country 对数据帧进行分组，然后应用一个函数来获取唯一计数，或者仅根据feature 的值获取计数，然后调用 to_frame 以从中创建帧。

(df.groupby('country', sort=False)
.apply(lambda x:pd.Series.nunique(x['ID']) 
                    if x['feature'].eq(2).any() 
                    else pd.Series.count(x['ID'])).to_frame('count')
 )

         count
country       
US           2
UK           2
Canada       2
Mexico       1

【讨论】：