Pandas两列值对比计数脚本运行结果错误如何修复

测试数据构造 #

import pandas as pd
d = {'a': [1, 3], 'b': [0, 2]}
df = pd.DataFrame(data=d, index=[1, 2])
print(df)

执行输出：

a  b
1  1  0
2  3  2

预期需求：逐行对比同位置a列和b列的数值，统计a大于b的次数、b大于a的次数，上述样例的预期返回结果为(2, 0)。

原有代码问题 #

原有实现使用双层嵌套循环，会把a列的每个值和b列的所有值做交叉对比（即笛卡尔积匹配），不是按行对齐同位置值比较，所以会返回错误结果(3, 1)。

正确实现 #

推荐方案：pandas向量化运算（性能最优） #

利用pandas Series的逐元素比较特性，直接对两列做大小判断，布尔结果中True等价于1、False等价于0，直接求和就能得到对应次数，不需要写任何循环，大数据量下性能远高于循环写法：

def diff(dataframe):
    a_counter = (dataframe["a"] > dataframe["b"]).sum()
    b_counter = (dataframe["b"] > dataframe["a"]).sum()
    return a_counter, b_counter

调用diff(df)即可得到正确结果(2, 0)。如果需要统计两列数值相等的次数，补充eq_counter = (dataframe["a"] == dataframe["b"]).sum()即可。

备选方案：逐行循环（仅适合极小数据集） #

如果一定要用循环实现，不要写双层嵌套，改为逐行迭代取同位置的两个值对比即可：

def diff(dataframe):
    a_counter = 0
    b_counter = 0
    for _, row in dataframe.iterrows():
        if row["a"] > row["b"]:
            a_counter += 1
        elif row["b"] > row["a"]:
            b_counter += 1
    return a_counter, b_counter

提示：数据集行数过万时，循环写法的运行速度会比向量化写法慢几十到上百倍，非特殊情况优先用向量化实现。

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