python - pandas 自定义 idxmax() 函数，具有特殊的相等性处理

Question

我正在处理一个充满选举数据的 csv 文件。我的原始样本可以表示为:

        city      party1     party2     party3
   0    city1     50         107        114
   1    city2     181        323        326
   2    city3     26         28         75
   3    city4     32         47         59
   4    ciy5      8          21         21

我使用 pandas 的 idxmax() 函数创建了一个名为“winner”的新列，如下所示:

 mydf['winner'] = mydf[['party1','party2','party3']].idxmax(axis=1)

我的目标是确定哪个政党在每个城市中处于第一位。结果如下:

        city      party1     party2     party3      winner
   0    city1     50         107        114         party3
   1    city2     181        323        326         party3
   2    city3     26         28         75          party3
   3    city4     32         47         59          party3
   4    ciy5      8          21         21          party2

最后一个raw的winner的值为false，因为party2和party3的分数相同。

是否可以在函数 idxmax 中包含一个异常(exception)，认为两个值相等并给出“相等”？

Answer 1

您可以使用 DataFrame.eq用于比较 subset 与 DataFrame.max每行的值，然后是 sum它们和哪里的值比 1 高，最多有重复项。那么可以用mask覆盖idxmax的值带掩码 s > 1:

a = mydf[['party1','party2','party3']]
mydf['winner'] = a.idxmax(axis=1)

s = a.eq(a.max(axis=1), axis=0).sum(axis=1)
print (s)
0    1
1    1
2    1
3    1
4    2
dtype: int64

mydf['winner'] = mydf['winner'].mask(s > 1, 'Equality')
print (mydf)
    city  party1  party2  party3    winner
0  city1      50     107     114    party3
1  city2     181     323     326    party3
2  city3      26      28      75    party3
3  city4      32      47      59    party3
4   ciy5       8      21      21  Equality

如果需要还可以通过 mul 的列值对多个 df 进行赋值, 然后 apply join 最后删除 , by strip :

a = mydf[['party1','party2','party3']]
df = a.eq(a.max(axis=1), axis=0)
print (df)
  party1 party2 party3
0  False  False   True
1  False  False   True
2  False  False   True
3  False  False   True
4  False   True   True

mydf['winner'] = df.mul(df.columns.to_series())
                   .apply(','.join, axis=1)
                   .str.strip(',')
print (mydf)
    city  party1  party2  party3         winner
0  city1      50     107     114         party3
1  city2     181     323     326         party3
2  city3      26      28      75         party3
3  city4      32      47      59         party3
4   ciy5       8      21      21  party2,party3