Numpy中高效对比单元格值按条件统计节点取值相等频率的方法

我们的目标是统计两个节点取值相等的出现频率，具体规则如下：

• 对于输入的二维数组，逐行遍历所有列对
• 若当前行内第k列和第p列的取值相等，则计数矩阵对应位置X[k,p]加1，否则不做累加
• 最终得到的计数矩阵就是所有行内列值相等的频率统计结果

示例说明 #

一维场景示例 #

首先以单个向量为例，输入数据：

pd.DataFrame([0,4,1,1,1],index=['A','B','C','D','E'])

对应输出：

0
A  0
B  4
C  1
D  1
E  1

按照规则生成的相等判断矩阵N如下：

A   B   C   D   E
A   1   0   0   0   0
B   0   1   0   0   0
C   0   0   1   1   1
D   0   0   1   1   1
E   0   0   1   1   1

二维场景示例 #

扩展到形状为(4,5)的二维数组，输入如下：

A  B  C  D  E
0  0  0  0  0  0
1  0  4  1  1  1
2  0  1  1  2  2
3  0  3  2  2  2

逐行计算相等判断矩阵后累加，最终得到的频率矩阵结果：

A    B    C    D    E
A  4.0  1.0  1.0  1.0  1.0
B  1.0  4.0  2.0  1.0  1.0
C  1.0  2.0  4.0  3.0  3.0
D  1.0  1.0  3.0  4.0  4.0
E  1.0  1.0  3.0  4.0  4.0

原有实现 #

目前已有的实现采用三层Python循环实现逻辑，代码如下：

import numpy as np
    
arr=[[ 0,0,0,0,0],
    [0,4,1,1,1],
    [0,1,1,2,2],
    [0,3,2,2,2]]
arr=np.array(arr)

# 行数
npart = len(arr[:,0])
# 列数
m = len(arr[0,:])
X = np.zeros(shape =(m,m), dtype = np.double)
for i in range(npart):
    for k in range(m):
        for p in range(m):
                # 判断当前行两列值是否相等
                if arr[i,k] == arr[i,p]:
                    X[k,p] = X[k,p] + 1

输出结果与示例完全一致。

可以直接使用Numpy的广播特性替代三层Python循环，所有运算在底层C层面执行，性能提升显著，尤其是数据量较大时优势更明显，优化后代码如下：

import numpy as np

arr = np.array([
    [0,0,0,0,0],
    [0,4,1,1,1],
    [0,1,1,2,2],
    [0,3,2,2,2]
])
# 广播对比所有列对的值，沿行维度求和得到计数矩阵
X = (arr[:, :, None] == arr[:, None, :]).sum(axis=0).astype(np.double)

运行后得到的结果与原有实现完全一致。

内容的提问来源于stack exchange，提问作者rpb