在一个数据帧中找到每行的前n个值,并使用这些索引在另一个数据帧中获取值以执行成对操作

时间:2017-09-19 18:54:34

标签: python pandas numpy dataframe

在这种情况下,我有两个数据帧A和B.

      c1  c2   c3              c1      c2     c3 
r0    7    6    4       r0      0       0      1
r1    6    2    5       r1      1       1      0
r2    3    5    9       r2      1       0      1

A是左边的数据框,右边是B.

基本上我的目标是在A的每一行中找到前2个值,在B中找到相应的行值,然后取这些对的乘积之和。

因此,例如在第一行中,A中的顶部值是7和6,它们对应于B的第一行中的0,0。然后我想返回7 * 0 + 6 * 0 = 0。我想在每一行都这样做,并返回类似的内容:

d1   0
d2   6
d3   9

我目前正在使用一个使用numpy argsort的实现来查找A的每一行中前n个值的索引,然后使用map和自定义函数来遍历行并找到product-sum

这种方法对我来说真的很慢,所以我想知道是否有更快的选择。谢谢。

1 个答案:

答案 0 :(得分:2)

使用rank获取前2个值并将其用作B的掩码。

In [1311]: (A*B.where(A.rank(axis=1) >= 2)).sum(axis=1)
Out[1311]:
r0    0.0
r1    6.0
r2    9.0
dtype: float64

详细

In [1314]: A.rank(axis=1)
Out[1314]:
     c1   c2   c3
r0  3.0  2.0  1.0
r1  3.0  1.0  2.0
r2  1.0  2.0  3.0

In [1315]: A.rank(axis=1) >=2
Out[1315]:
       c1     c2     c3
r0   True   True  False
r1   True  False   True
r2  False   True   True

In [1317]: B.where(A.rank(axis=1) >= 2)
Out[1317]:
     c1   c2   c3
r0  0.0  0.0  NaN
r1  1.0  NaN  0.0
r2  NaN  0.0  1.0

In [1318]: (A*B.where(A.rank(axis=1) >= 2))
Out[1318]:
     c1   c2   c3
r0  0.0  0.0  NaN
r1  6.0  NaN  0.0
r2  NaN  0.0  9.0