python - numpy 中的元素运算

Question

我有一个4000 x 16矩阵A。我有 256 x 1 向量 B。我需要对 A 中的 each 元素与 B 的 every 元素进行逐元素加法，并获得一个 3D 数组尺寸 4000 x 16 x 256。在没有 numpy 循环的情况下实现这一目标的最有效方法是什么？

Answer 1

可以先将A整形为4000×16×1矩阵，将B整形为1×1×256矩阵。然后你可以执行加法:

A[:,:,None] + B.reshape(1, 1, -1)

对于 A 一个 4×5 矩阵，B 一个 3×1 矩阵，例如，我们得到:

>>> A
array([[1, 2, 2, 2, 2],
       [2, 1, 0, 3, 2],
       [4, 1, 2, 1, 3],
       [3, 2, 3, 2, 2]])
>>> B
array([[0],
       [3],
       [3],
       [6],
       [4],
       [3]])
>>> A[:,:,None] + B.reshape(1, 1, -1)
array([[[ 1,  4,  4,  7,  5,  4],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5]],

       [[ 2,  5,  5,  8,  6,  5],
        [ 1,  4,  4,  7,  5,  4],
        [ 0,  3,  3,  6,  4,  3],
        [ 3,  6,  6,  9,  7,  6],
        [ 2,  5,  5,  8,  6,  5]],

       [[ 4,  7,  7, 10,  8,  7],
        [ 1,  4,  4,  7,  5,  4],
        [ 2,  5,  5,  8,  6,  5],
        [ 1,  4,  4,  7,  5,  4],
        [ 3,  6,  6,  9,  7,  6]],

       [[ 3,  6,  6,  9,  7,  6],
        [ 2,  5,  5,  8,  6,  5],
        [ 3,  6,  6,  9,  7,  6],
        [ 2,  5,  5,  8,  6,  5],
        [ 2,  5,  5,  8,  6,  5]]])

性能:如果我使用 A 一个 4000×16 矩阵和 B 一个 256×1 float 矩阵运行上述 100 次点，我得到以下结果:

>>> timeit(lambda: A[:,:,None] + B.reshape(1, 1, -1), number=100)
5.949456596048549

因此单次运行大约需要 59.49457 毫秒。计算 16'384'000 个元素看起来很合理。