Python / numpy：如何沿轴以指定长度替换矩阵元素

Question

我有一个数组，该数组指定要更改为0的矩阵的最后一个轴上的元素数。如何有效地做到这一点？

x = np.ones((4, 4, 10))
change_to_zeros = np.random.randint(10, size=(4, 4))
# change_to_zeros = 
# [[2 1 6 8]
# [4 0 4 8]
# [7 6 6 2]
# [4 0 7 1]]

现在我想做的是类似x[:, :, :change_to_zeros] = 0的事情-例如，对于change_to_zeros的第一个元素，我有change_to_zeros[0, 0] = 2，所以我想更改第一个（或最后一个，或沿x的最后一个轴（长度10）为0的2个元素。

说明：例如，在x[0, 0, :]，我有长度为10的那些。我想将其中的两个（change_to_zeros[0, 0] = 2）更改为0，其余的保留为1。

Answer 1

您可以使用x创建一个布尔数组（形状与change_to_zeros[:,:,None] > np.arange(x.shape[-1]) 相同），然后将零赋给 true s：

x[change_to_zeros[:,:,None] > np.arange(x.shape[-1])] = 0

检查结果：

change_to_zeros[0,0]
# 2

x[0,0]
# array([ 0.,  0.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.])

change_to_zeros[0,2]
# 7

x[0,2]
# array([ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  1.,  1.])