如何用向量化方式对NumPy二维数组行内非零元素做均值减法归一化

需求说明 #

在推荐系统场景中，我们用NumPy数组存储用户评分数据：每行代表一个用户，每列代表一部电影。需要对每行做归一化处理：仅对行内的非零评分条目，减去该行所有非零评分的均值，零值条目保持不变。

示例输入 #

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

目标效果（循环实现） #

以下循环代码可以实现需求，但我们需要更高效的向量化替代方案：

for i in range(len(x)):
    y = np.mean(x[i][x[i].nonzero()])
    x[i][x[i].nonzero()] -= y

输出结果：

[[ 0.  0. -1.  0.  1.]
 [ 0.  0. -1.  0.  1.]
 [ 0.  0. -1.  0.  1.]]

问题分析 #

你之前尝试的代码使用了整行均值（包含零值），不符合需求：

mask = x!= 0
t = np.where(mask, x-x.mean(axis=1).reshape(-1,1), x)

我们需要计算的是每行仅非零元素的均值，而非整行均值。

向量化解决方案 #

通过以下步骤实现无循环的向量化处理：

1. 生成非零元素的掩码矩阵
2. 计算每行非零元素的均值
3. 仅在非零位置应用均值减法，零位置保持原值

完整代码：

import numpy as np

x = np.array([[0,0,1,2,3],[0,0,2,3,4],[0,0,3,4,5.0]])
mask = x != 0

# 计算每行非零元素的均值：非零元素总和 / 非零元素数量
row_nonzero_means = (x * mask).sum(axis=1) / mask.sum(axis=1)

# 将均值转为列向量，与原矩阵做减法，仅在非零位置生效
normalized_x = np.where(mask, x - row_nonzero_means.reshape(-1, 1), x)

print(normalized_x)

输出结果与循环实现完全一致：

[[ 0.  0. -1.  0.  1.]
 [ 0.  0. -1.  0.  1.]
 [ 0.  0. -1.  0.  1.]]

代码解释 #

• mask = x != 0：创建布尔矩阵，标记所有非零评分的位置
• (x * mask).sum(axis=1)：计算每行非零元素的总和（零值被mask置为0，不影响求和）
• mask.sum(axis=1)：统计每行非零元素的数量
• row_nonzero_means.reshape(-1, 1)：将一维均值数组转为列向量，保证与原矩阵维度匹配，支持广播运算
• np.where(mask, ...)：仅在非零位置执行均值减法，零位置保留原值

内容的提问来源于stack exchange，提问作者Washik Uddin Ahmed Mollah