如何同时切入阵列的两个轴？

如何解决如何同时切入阵列的两个轴？

问题： x 是形状为 (100,50,2) 的 n 维数组。我想以矢量化的方式修改在第一和第二轴具有相同索引的数组元素。 没有矢量化，我可以做到：

for i in range(50):
   x[:,i,:] = 0

问题与访问由第一和第二轴定义的对角线元素相同。

---

此问题已在 Modifying diagonals in multidimensional numpy arrays 中提出。但是，那里唯一的答案是使用函数 numpy.lib.stride_tricks。根据 numpy 文档，此函数“可能指向无效内存并可能破坏结果或使程序崩溃”和“建议尽可能避免 as_strided”。由于这些原因，我想知道是否有更安全、更简单的方法来解决这个问题。

解决方法

您可以使用 np.arange(50) 作为两个位置的索引器；它将给出对 (0,0),(1,1),...,(49,49) 作为这两个位置的索引器：

n = np.arange(x.shape[2])
x[:,n,:] = 0

---

示例：

>>> a = np.arange(150).reshape(3,5,2)
>>> n = np.arange(a.shape[2])
>>> a[:,:] = 0
>>> a

array([[[[  0,0],# changed
         [  2,3],[  4,5],[  6,7],[  8,9]],[[ 10,11],[  0,# changed
         [ 14,15],[ 16,17],[ 18,19]],[[ 20,21],[ 22,23],# changed
         [ 26,27],[ 28,29]],[[ 30,31],[ 32,33],[ 34,35],# changed
         [ 38,39]],[[ 40,41],[ 42,43],[ 44,45],[ 46,47],0]]],# changed


       [[[  0,# changed
         [ 52,53],[ 54,55],[ 56,57],[ 58,59]],[[ 60,61],# changed
         [ 64,65],[ 66,67],[ 68,69]],[[ 70,71],[ 72,73],# changed
         [ 76,77],[ 78,79]],[[ 80,81],[ 82,83],[ 84,85],# changed
         [ 88,89]],[[ 90,91],[ 92,93],[ 94,95],[ 96,97],# changed
         [102,103],[104,105],[106,107],[108,109]],[[110,111],# changed
         [114,115],[116,117],[118,119]],[[120,121],[122,123],# changed
         [126,127],[128,129]],[[130,131],[132,133],[134,135],# changed
         [138,139]],[[140,141],[142,143],[144,145],[146,147],0]]]])            # changed