a = np.array([[1, 2], [1, 2]])
print(a.shape)
# (2, 2)
# indexing with np.newaxis inserts a new 3rd dimension, which we then repeat the
# array along, (you can achieve the same effect by indexing with None, see below)
b = np.repeat(a[:, :, np.newaxis], 3, axis=2)
print(b.shape)
# (2, 2, 3)
print(b[:, :, 0])
# [[1 2]
# [1 2]]
print(b[:, :, 1])
# [[1 2]
# [1 2]]
print(b[:, :, 2])
# [[1 2]
# [1 2]]
话虽如此,您通常可以通过使用broadcast避免完全重复阵列。例如,假设我要添加一个(3,)
向量:
c = np.array([1, 2, 3])
到a
。我可以a
在第三维中复制3次的内容,然后c
在第一维和第二维中复制2次的内容,这样我的两个数组都是(2, 2, 3)
,然后计算它们的总和。但是,这样做更加简单快捷:
d = a[..., None] + c[None, None, :]
在此,a[..., None]
具有形状,(2, 2, 1)
并且c[None, None, :]
具有形状(1, 1, 3)
*。当我计算总和时,结果沿大小为1的维度“广播”出去,给了我shape的结果(2, 2, 3)
:
print(d.shape)
# (2, 2, 3)
print(d[..., 0]) # a + c[0]
# [[2 3]
# [2 3]]
print(d[..., 1]) # a + c[1]
# [[3 4]
# [3 4]]
print(d[..., 2]) # a + c[2]
# [[4 5]
# [4 5]]
广播是一项非常强大的技术,因为它避免了在内存中创建输入数组的重复副本所涉及的额外开销。
*尽管为清楚起见,我将它们包括在内,但实际上并不需要None
索引c
-您也可以这样做a[..., None] + c
,即(2, 2, 1)
针对(3,)
数组广播数组。这是因为,如果其中一个数组的尺寸小于另一个数组的尺寸,则仅两个数组的 尾部 尺寸需要兼容。举一个更复杂的例子:
a = np.ones((6, 1, 4, 3, 1)) # 6 x 1 x 4 x 3 x 1
b = np.ones((5, 1, 3, 2)) # 5 x 1 x 3 x 2
result = a + b # 6 x 5 x 4 x 3 x 2