如何进一步提升tArray数据归一化代码的时间效率？

我有一个形状为(11, 512)的矩阵tArray，需要对其中的值进行归一化处理。发现np.max()函数耗时较高，但未找到进一步优化的方法，能否优化以下代码的时间效率？

原代码：

tArray = np.array([[val/tArray[i][sqLen-1] for val in tArray[i]] if i not in [1,2] else [val/np.max(tArray[i][:sqLen-1]) for val in tArray[i]] for i in range(len(tArray))])

复现代码：

tArray = np.random.randint(1, 100, size=(11, 512))
tArray = np.array([[val/tArray[i][512-1] for val in tArray[i]] if i not in [1,2] else [val/np.max(tArray[i][:512-1]) for val in tArray[i]] for i in range(len(tArray))])

问题分析 #

原代码的核心低效点在于：

• 使用Python层面的列表推导逐元素循环，远慢于NumPy底层C实现的向量化操作；
• 对索引1、2的行，在循环中重复调用np.max()并逐元素计算，额外增加了大量不必要的开销。

优化方案 #

改用NumPy的向量化操作+广播机制，一次性完成所有行的归一化计算，彻底避免Python循环：

import numpy as np

# 生成测试数据
tArray = np.random.randint(1, 100, size=(11, 512))

# 1. 初始化除数数组：默认取每行最后一个元素
divisors = tArray[:, -1].copy()

# 2. 对索引1、2的行，替换为该行前511个元素的最大值
divisors[[1, 2]] = np.max(tArray[[1, 2], :-1], axis=1)

# 3. 利用广播机制完成整矩阵的归一化
tArray_normalized = tArray / divisors[:, np.newaxis]

优化效果说明 #

• 仅对需要的两行调用一次np.max()，避免重复计算；
• 广播除法让所有元素的计算在NumPy底层批量完成，效率比原代码提升数十倍甚至上百倍；
• 代码结构更简洁，可读性更强。

时间对比示例 #

用timeit测试1000次运行的耗时（仅供参考，具体数值取决于硬件）：

import timeit

def original_code():
    tArray = np.random.randint(1, 100, size=(11, 512))
    tArray = np.array([[val/tArray[i][512-1] for val in tArray[i]] if i not in [1,2] else [val/np.max(tArray[i][:512-1]) for val in tArray[i]] for i in range(len(tArray))])
    return tArray

def optimized_code():
    tArray = np.random.randint(1, 100, size=(11, 512))
    divisors = tArray[:, -1].copy()
    divisors[[1, 2]] = np.max(tArray[[1, 2], :-1], axis=1)
    tArray_normalized = tArray / divisors[:, np.newaxis]
    return tArray_normalized

print("原代码耗时：", round(timeit.timeit(original_code, number=1000), 4), "秒")
print("优化后代码耗时：", round(timeit.timeit(optimized_code, number=1000), 4), "秒")

输出示例：

原代码耗时： 1.2345 秒
优化后代码耗时： 0.0123 秒

注意事项 #

如果矩阵中存在0值，需要提前处理避免除以0的情况，比如给除数加一个极小的常数（如1e-8）。

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