按照两个列的条件聚合DataFrame的pandas组

Question

我得到了这个样本数据框：

df = pd.DataFrame({'CUSTOM_CRITERIA':[1111,22222,1111,1212,1212,3333,5555, 1111], 
                'AD_UNIT_NAME':['inp2_l_d', 'inp1', 'pixel_d', 'inp2_l_d', 'anchor_m','anchor_m','anchor_m','inp2_l_d'], 
                'TOTAL_CODE_SERVED_COUNT':[10, 20, 10, 12, 18,500,100,50]})

我需要针对每个CUSTOM_CRITERIA通过条件获取最大的total_code_served_count，该条件为：已提供代码更多 -> anchor_m [total_code served] OR inp2_l_d[total_code served] + pixel_d[total_code served]。

我的当前解决方案如下：

data_dict = clean_data.to_dict(orient='records')

for item in data_dict:
    desktop_impression_max_calculated = sum([d['TOTAL_CODE_SERVED_COUNT'] for d in data_dict if d['CUSTOM_CRITERIA'] == item['CUSTOM_CRITERIA'] and ('inp2_l_d' in d['AD_UNIT_NAME'].lower() or 'pixel_d' in d['AD_UNIT_NAME'].lower())])
    mobile_impression_max_calculated = sum([d['TOTAL_CODE_SERVED_COUNT'] for d in data_dict if d['CUSTOM_CRITERIA'] == item['CUSTOM_CRITERIA'] and 'anchor_m' in d['AD_UNIT_NAME'].lower()])
    item['IMPRESSIONS_MAX'] = max(desktop_impression_max_calculated,mobile_impression_max_calculated)

clean_data = pd.DataFrame(data_dict)   
agg_map = {'IMPRESSIONS_MAX': 'first' }

clean_data = clean_data.groupby('CUSTOM_CRITERIA').agg(agg_map).reset_index()

当存在大量数据时，由于N^2复杂度，这需要很长时间才能运行。我相信使用pandas可以找到更好、更简单的方法。

- Maxim Kogan

Answer 1

您可以通过将TOTAL_CODE_SERVED_COUNT列中的值乘以布尔掩码m1和m2来创建两个掩码列，然后在CUSTOM_CRITERIA上对这些掩码列进行分组并使用sum进行聚合，最后沿着axis=1取max来获得最终结果：

m1 = df['AD_UNIT_NAME'].str.contains(r'(?i)inp2_l_d|pixel_d')
m2 = df['AD_UNIT_NAME'].str.contains(r'(?i)anchor_m')

pd.DataFrame((df['TOTAL_CODE_SERVED_COUNT'].values * [m1, m2]).T)\
  .groupby(df['CUSTOM_CRITERIA']).sum().max(1).reset_index(name='IMPRESSIONS_MAX')

---

   CUSTOM_CRITERIA  IMPRESSIONS_MAX
0             1111               70
1             1212               18
2             3333              500
3             5555              100
4            22222                0

- Shubham Sharma