pandas行级nlargest(2)取列名时如何处理等值并列与全0行问题

实现逻辑 #

• 逐行先过滤值为0的国家数据，0值不参与排名，解决全0行返回错误列的问题
• 对非0数据取第2高的访问量作为阈值，所有大于等于阈值的国家全部保留，解决并列值被截断的问题
• 结果长度统一补0对齐到3列，和期望输出格式一致

可直接运行的代码 #

import pandas as pd

# 原始示例DataFrame
df = pd.DataFrame({
    'Name': ['Sara', 'Adam','Ciara', 'John','Paul'],
    'UK': [1,2,0,4,2],
    'US': [1,3,0,1,2],
    'IN': [1,5,0,1,0],
    'CA': [0,0,0,7,1],
    'JP': [0,0,0,0,1]
})

country_cols = ['UK','US','IN','CA','JP']
io = df[country_cols]

def extract_top2(row):
    # 过滤0值，排除无访问量的国家
    valid_data = row[row > 0]
    # 全0行直接返回占位值
    if valid_data.empty:
        return [0, 0, 0]
    # 取Top2的临界值，所有大于等于该值的国家都纳入结果
    threshold = valid_data.nlargest(2).iloc[-1]
    res = valid_data[valid_data >= threshold].index.to_list()
    # 补0对齐到3列
    res += [0] * (3 - len(res))
    return res

# 生成结果并拼接原表
top_res = io.apply(extract_top2, axis=1, result_type='expand')
top_res.columns = [0, 1, 2]
final_df = top_res.join(df)

运行输出结果 #

0   1   2  Name  UK  US  IN  CA  JP
0  UK  US  IN  Sara   1   1   1   0   0
1  IN  US   0  Adam   2   3   5   0   0
2   0   0   0 Ciara   0   0   0   0   0
3  CA  UK   0  John   4   1   1   7   0
4  UK  US   0  Paul   2   2   0   1   1

原写法的问题说明 #

• 直接调用nlargest(2)默认固定返回2个元素，遇到并列同值场景会按列顺序截断，不会返回所有达到Top2阈值的列
• 没有提前过滤0值，全0行计算时会直接取前两列的0作为结果，导致返回错误的国家名

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