python - 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）的第一个（或最后一个，或其他）2 个元素更改为 0。

澄清：例如，x[0, 0, :]我有长度为 10 的。我想将其中的 2 个 ( change_to_zeros[0, 0] = 2) 更改为 0，其余的保留为 1。

标签： pythonnumpy

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.])