基于元素归一化的大型矩阵行内高值计数优化需求

我之前在StackOverflow询问过矩阵行/列内高值求和问题，得到的方案无法支持元素级操作。

输入的相关矩阵如下：

array([[-1,  7, -2,  1,  4],
       [ 6,  3,  -3,  5,  1]])

需求说明：

• 对矩阵中每个元素先做归一化处理：元素为负则加3，为正则减3
• 统计每个元素所在行中，比该元素处理后的值大的元素数量

举个例子：

• 第一行第一个元素为-1（负），加3后该行变为2 7 -2 1 4，比2大的元素有7和4，数量为2
• 第一行第二个元素为7（正），减3后该行变为-1 4 -2 1 4，比4大的元素数量为0

期望的行处理输出结果：

2 0 2 3 1 
1 3 4 2 3

由于矩阵规模为11k×12k，循环实现耗时过长，因此需要矢量化方案。当前尝试的vectorize+rankdata方法会对整行统一处理，无法实现元素级的单独计算，代码如下：

corr = np.array([[-1,  7, -2,  1,  4],
                 [ 6,  3,  -3,  5,  1]])


def element_wise_bigger_than(x, axis):
    return x.shape[axis] - rankdata(x, method='max', axis=axis)


ld = lambda t: t + 3 if t < 0 else t - 3
f = np.vectorize(ld)


element_wise_bigger_than(f(corr), 1)

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