pandas idxmax：return 所有行，如果是平局

Question

我正在使用一个数据框，我在其中按概率对每一行进行加权。现在，我想 select 概率最高的行，我正在使用 pandas idxmax() 这样做，但是当有关系时，它只是 returns 中的第一行那些打领带的。在我的例子中，我想获取所有匹配 的行。

此外，我这样做是作为一个研究项目的一部分，我正在处理数百万个数据帧，如下所示，因此保持快速是一个问题。

示例：

我的数据是这样的：

data = [['chr1',100,200,0.2],
    ['ch1',300,500,0.3],
    ['chr1', 300, 500, 0.3],
    ['chr1', 600, 800, 0.3]]

从这个列表中，我创建了一个 pandas 数据框，如下所示：

weighted = pd.DataFrame.from_records(data,columns=['chrom','start','end','probability'])

看起来像这样：

  chrom  start  end  probability
0  chr1    100  200          0.2
1   ch1    300  500          0.3
2  chr1    300  500          0.3
3  chr1    600  800          0.3

然后 select 符合 argmax（概率）的行使用：

selected =  weighted.ix[weighted['probability'].idxmax()]

当然是returns:

chrom          ch1
start          300
end            500
probability    0.3
Name: 1, dtype: object

是否有一种（快速）方法可以在存在关系时获取所有值？

谢谢！

Answer 1

嗯，这可能是您正在寻找的解决方案：

weighted.loc[weighted['probability']==weighted['probability'].max()].T
#               1     2     3
#chrom        ch1  chr1  chr1
#start        300   300   600
#end          500   500   800
#probability  0.3   0.3   0.3

Answer 2

瓶颈在于计算布尔索引器。您可以通过使用底层 NumPy 数组执行计算来绕过与 pd.Series 对象相关的开销：

df2 = df[df['probability'].values == df['probability'].values.max()]

与 Pandas 等效的性能基准测试：

# tested on Pandas v0.19.2, Python 3.6.0

df = pd.concat([df]*100000, ignore_index=True)

%timeit df['probability'].eq(df['probability'].max())               # 3.78 ms per loop
%timeit df['probability'].values == df['probability'].values.max()  # 416 µs per loop

pandas idxmax：return 所有行，如果是平局

pandas idxmax: return all rows in case of ties

python

pandas

performance

argmax