熊猫idxmax:如果有联系,则返回所有行

时间:2018-10-01 09:39:43

标签: python pandas performance argmax

我正在使用一个数据框,在该数据框中我按行的权重对其进行加权。现在,我想选择概率最高的行,并且正在使用pandas idxmax()进行此操作,但是当有联系时,它只返回联系中的第一行。就我而言,我想获取所有相关的行

此外,作为研究项目的一部分,我正在执行此操作,在该项目中,我正在处理数以百万计的数据帧,如下所示,因此保持快速是一个问题。

示例:

我的数据如下:

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

从此列表中,我创建了一个熊猫数据框,如下所示:

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

然后使用以下命令选择适合argmax(probability)的行:

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

当然会返回:

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

有关系时,是否有(快速)获取所有值的方法?

谢谢!

2 个答案:

答案 0 :(得分:3)

瓶颈在于计算布尔索引器。您可以通过对基础的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

答案 1 :(得分:2)

好吧,这可能是您正在寻找的解决方案:

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