面试题：Python中NumPy多维数组运算及性能优化

Question

假设有两个形状为(1000, 1000)的二维NumPy数组A和B，要求计算A * B（对应元素相乘），然后对结果矩阵的每一行求平均值。请写出Python代码，并说明如何通过NumPy的特性对代码进行性能优化，以提高运算效率。

Answer 1

import numpy as np

# 生成两个形状为(1000, 1000)的二维NumPy数组
A = np.random.rand(1000, 1000)
B = np.random.rand(1000, 1000)

# 对应元素相乘并对结果矩阵的每一行求平均值
result = np.mean(A * B, axis=1)
print(result)

性能优化说明

1. 向量化操作：在NumPy中，使用A * B进行对应元素相乘，而不是使用Python的循环语句。NumPy的底层实现是基于C语言，向量化操作利用了这种底层优化，大大减少了Python循环带来的额外开销。
2. np.mean函数：使用np.mean函数并指定axis=1来计算每一行的平均值，np.mean同样是经过优化的向量化操作，避免了手动编写循环来计算平均值，提高了运算效率。