python - 在给定轴索引处替换 nd numpy 数组中的值

Question

我想在给定索引(例如 0)沿给定轴(例如 - 2) 具有给定值(例如 0)，等效于:

A[:,:,0,:]=0

问题是输入数组 A 可能以 3D 或 4D 或其他形状出现，所以对于 3D 数据我需要

A[:,0,:]=0

如果是 5D:A[:,:,:,0,:]=0

目前我正在使用 exec() 来完成这项工作:

slicestr=[':']*numpy.ndim(var)
slicestr[-2]=str(0)
slicestr=','.join(slicestr)
cmd='A[%s]=0' %slicestr
exec(cmd)
return A

我有点担心 exec() 的使用可能不是一个很好的方法。我知道 numpy.take() 可以为我提供沿特定轴的特定索引处的列，但要替换值，我仍然需要构建动态的切片/索引字符串。所以我想知道是否有任何原生的 numpy 方法可以实现这一目标？

谢谢。

Answer 1

您可以使用省略号(有关更多信息，请参阅 numpy indexing)跳过前几个维度:

# assert A.ndim >= 2
A[...,0,:]=0

---

A2d = np.arange(12).reshape(2,6)
A3d = np.arange(12).reshape(2,3,2)
A4d = np.arange(12).reshape(2,1,3,2)

A2d[...,0,:] = 0

A2d
#array([[ 0,  0,  0,  0,  0,  0],
#       [ 6,  7,  8,  9, 10, 11]])

A3d[...,0,:] = 0

A3d
#array([[[ 0,  0],
#        [ 2,  3],
#        [ 4,  5]],

#       [[ 0,  0],
#        [ 8,  9],
#        [10, 11]]])

A4d[...,0,:] = 0

A4d
#array([[[[ 0,  0],
#         [ 2,  3],
#         [ 4,  5]]],  

#       [[[ 0,  0],
#         [ 8,  9],
#         [10, 11]]]])