数据帧行的快速笛卡尔求和

Fast Cartesian sum of rows of dataframe

我在 3 轴 (x, y, z) 上有两个错误数据帧:

df1 = pd.DataFrame([[0, 1, 2], [-1, 0, 1], [-2, 0, 3]], columns = ['x', 'y', 'z'])
df2 = pd.DataFrame([[1, 1, 3], [1, 0, 2], [1, 0, 3]], columns = ['x', 'y', 'z'])

我正在寻找一种快速方法来计算两个数据帧每一行的笛卡尔平方和。

编辑 我目前的解决方案:

cartesian_sum = list(np.sum(list(tup), axis = 0).tolist() 
                    for tup in itertools.product( (df1**2).to_numpy().tolist(),
                                                  (df2**2).to_numpy().tolist() ) )

cartesian_sum
>>> 
[[1, 2, 13],
 [1, 1, 8],
 [1, 1, 13],
 [2, 1, 10],
 [2, 0, 5],
 [2, 0, 10],
 [5, 1, 18],
 [5, 0, 13],
 [5, 0, 18]]

太慢(~ 2.4 毫秒;与完全基于 Pandas 运行 ~ 8-10 毫秒的解决方案相比)。

这与 related question (link here) 类似,但使用 itertools 太慢了。在 Python 中有更快的方法吗?

我认为您首先需要交叉连接,删除列 a,平方,将列转换为 MultiIndex 和每个第一级的 sum

df = df1.assign(a=1).merge(df2.assign(a=1), on='a').drop('a', axis=1) ** 2
df.columns = df.columns.str.split('_', expand=True)
df = df.sum(level=0, axis=1)
print (df)
   x  y   z
0  1  2  13
1  1  1   8
2  1  1  13
3  2  1  10
4  2  0   5
5  2  0  10
6  5  1  18
7  5  0  13
8  5  0  18

详情:

print (df1.assign(a=1).merge(df2.assign(a=1), on='a'))
   x_x  y_x  z_x  a  x_y  y_y  z_y
0    0    1    2  1    1    1    3
1    0    1    2  1    1    0    2
2    0    1    2  1    1    0    3
3   -1    0    1  1    1    1    3
4   -1    0    1  1    1    0    2
5   -1    0    1  1    1    0    3
6   -2    0    3  1    1    1    3
7   -2    0    3  1    1    0    2
8   -2    0    3  1    1    0    3

提高性能的一个想法:

#
def cartesian_product_simplified_changed(left, right):
    la, lb = len(left), len(right)
    ia2, ib2 = np.broadcast_arrays(*np.ogrid[:la,:lb])

    a = np.column_stack([left.values[ia2.ravel()] ** 2, right.values[ib2.ravel()] ** 2])
    a = a[:, :la] + a[:, la:]
    return a


a = cartesian_product_simplified_changed(df1, df2)
print (a)
[[ 1  2 13]
 [ 1  1  8]
 [ 1  1 13]
 [ 2  1 10]
 [ 2  0  5]
 [ 2  0 10]
 [ 5  1 18]
 [ 5  0 13]
 [ 5  0 18]]