面试题：Python中NumPy广播机制在复杂数组运算中的应用

Question

现有两个NumPy数组，数组A的形状是(2, 3, 1)，数组B的形状是(4, 1, 3)，如果要让这两个数组进行元素级乘法（* 运算符），请说明广播机制是如何起作用的，并且编写代码实现这个运算。

Answer 1

1. 广播机制说明：
   • 数组A形状为(2, 3, 1)，数组B形状为(4, 1, 3)。
   • 从后往前比较维度，最后一维A是1，B是3，不相等，但A维度为1，可以广播扩展成和B最后一维一样的3。
   • 倒数第二维，A是3，B是1，不相等，但B维度为1，可以广播扩展成和A倒数第二维一样的3。
   • 第一维，A是2，B是4，不相等且都不为1，不能直接广播。需要在A的前面增加一个维度，使其形状变为(1, 2, 3, 1)，然后广播扩展成(4, 2, 3, 3)，B广播扩展成(4, 2, 3, 3)，这样两个数组形状相同就可以进行元素级乘法。
2. 代码实现：

import numpy as np

A = np.random.rand(2, 3, 1)
B = np.random.rand(4, 1, 3)
result = A * B
print(result)