0

我正在尝试用 numpy 数组实现一个简单的算法。下面是我想要完成的一个例子。

设 a 为 4x3 形状的 numpy 数组

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

设 b 为形状为 2x3 的布尔数组

[[ True  False False]
 [ False False  True]]

b[0,0] 为真。所以我想要 a[0:2,0] = 100

b[1,2] 为真。所以我想要 a[2:4,2] = 100

即如果 b[row,col] 为真，我想要a[2*row:2*(row+1),col] = 100

什么有效：

  1. 使用虚拟轴将 b 的形状从 (2x3) 更改为 (2,1,3)
  2. 将 a 从 (4x3) 重塑为 (2,2,3)。我们称之为 a_reshape
  3. a_reshape[:,0:1,:][b] = 100
  4. a_reshape[:,1:2,:][b] = 100

什么不起作用

如果我尝试一次更改两列，即

a_reshape[:,0:2,:][b] = 100

我收到以下错误：

"boolean index did not match indexed array along dimension 1; dimension is 2 but corresponding boolean dimension is 1"

我可以添加a和b。在这种情况下，广播工作。但是我无法在索引时进行广播。

如果有人能解释为什么会这样，我将不胜感激。还有一种更好（更快）的方法来实现上述代码吗？任何提示将不胜感激。

谢谢

"""

import numpy as np

# a is of shape 4x3
a = np.arange(0,12).reshape(4,3)

print("Original array is \n")
print(a,'\n')

# b is of shape (2,1,3)
b = np.random.randint(0,2,(2,3)) == 0

print("Indexing array is \n")
print(b,'\n')

# Change shape of b from 2x3 to 2x1x3
b = b[:,None,:]

# a_reshape is of shape (2,2,3)
a_reshape = a.reshape(b.shape[0],-1,a.shape[-1])

# This works
a_reshape[:,0:1,:][b] = 100
a_reshape[:,1:2,:][b] = 100
print("New array is \n")
print(a)
print("Works")

# This does not. Gives error
# boolean index did not match indexed array along dimension 1; dimension is 2 but corresponding boolean dimension is 1
print("Trying to broadcast at once array is \n")
a_reshape[:,0:2,:][b] = 100
4

0 回答 0