如何在Python Pandas中基于列最大值生成TOP_COUNT与TOP_SUM新列

原始数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'ID': [111, 222, 333, 444, 555, 666],
    'COUNT_COL_A': [10, 15, 0, 20, 0, 10],
    'COUNT_COL_B': [10, 80, 0, 5, 0, 20],
    'SUM_COL_A': [320, 500, 110, 0, 0, 60],
    'SUM_COL_B': [120, 500, 350, 0, 0, 50]
})

需求 #

• 创建新列TOP_COUNT：
  1. 存储COUNT_COL_A或COUNT_COL_B中值更大的列名；
  2. 若两COUNT列值相等，则取对应SUM列值更大的那个COUNT列名；
  3. 若两COUNT列均为0，则TOP_COUNT为NaN。
• 创建新列TOP_SUM：
  1. 存储SUM_COL_A或SUM_COL_B中值更大的列名；
  2. 若两SUM列值相等，则取对应COUNT列值更大的那个SUM列名；
  3. 若两SUM列均为0，则TOP_SUM为NaN。

期望输出 #

ID   | COUNT_COL_A | COUNT_COL_B | SUM_COL_A | SUM_COL_B  | TOP_COUNT   | TOP_SUM
-----|-------------|-------------|-----------|------------|-------------|--------- 
111  | 10          | 10          | 320       | 120        | COUNT_COL_A | SUM_COL_A 
222  | 15          | 80          | 500       | 500        | COUNT_COL_B | SUM_COL_B  
333  | 0           | 0           | 110       | 350        | NaN         | SUM_COL_B  
444  | 20          | 5           | 0         | 0          | COUNT_COL_A | NaN
555  | 0           | 0           | 0         | 0          | NaN         | NaN
666  | 10          | 20          | 60        | 50         | COUNT_COL_B | SUM_COL_A

实现代码 #

处理TOP_COUNT列 #

利用numpy.where实现多层条件判断，保证矢量化运算效率（适合大数据量场景）：

# 判断两COUNT列是否均为0
mask_count_zero = (df['COUNT_COL_A'] == 0) & (df['COUNT_COL_B'] == 0)
# COUNT列不等时，直接取数值更大的列名
count_compare = np.where(df['COUNT_COL_A'] > df['COUNT_COL_B'], 'COUNT_COL_A', 'COUNT_COL_B')
# COUNT列相等时，根据对应SUM列的大小选择COUNT列名
sum_compare_for_count = np.where(df['SUM_COL_A'] > df['SUM_COL_B'], 'COUNT_COL_A', 'COUNT_COL_B')
# 合并所有逻辑生成TOP_COUNT
df['TOP_COUNT'] = np.where(
    mask_count_zero, 
    np.nan, 
    np.where(df['COUNT_COL_A'] == df['COUNT_COL_B'], sum_compare_for_count, count_compare)
)

处理TOP_SUM列 #

同理，按SUM列优先级+COUNT列兜底的逻辑实现：

# 判断两SUM列是否均为0
mask_sum_zero = (df['SUM_COL_A'] == 0) & (df['SUM_COL_B'] == 0)
# SUM列不等时，直接取数值更大的列名
sum_compare = np.where(df['SUM_COL_A'] > df['SUM_COL_B'], 'SUM_COL_A', 'SUM_COL_B')
# SUM列相等时，根据对应COUNT列的大小选择SUM列名
count_compare_for_sum = np.where(df['COUNT_COL_A'] > df['COUNT_COL_B'], 'SUM_COL_A', 'SUM_COL_B')
# 合并所有逻辑生成TOP_SUM
df['TOP_SUM'] = np.where(
    mask_sum_zero, 
    np.nan, 
    np.where(df['SUM_COL_A'] == df['SUM_COL_B'], count_compare_for_sum, sum_compare)
)

验证结果 #

执行上述代码后，打印df即可得到与期望一致的输出：

print(df.to_string(index=False))

内容的提问来源于stack exchange，提问作者dingaro