How to slice through two axis of an array simultaneously?

Question

Problem: x is an n-dimensional array with shape (100, 50, 50, 2). I would like to modify the elements of the array that have the same index in the 1st and 2nd axis in a vectorized way. Without vectorizing, I could do it with:

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

The problem is the same as accessing the elements of the diagonal defined by the 1st and 2nd axis.

---

This question was already asked in Modifying diagonals in multidimensional numpy arrays. However, the only answer there uses the function numpy.lib.stride_tricks. According to numpy documentation, this function "can point to invalid memory and can corrupt results or crash your program" and "it is advisable to avoid as_strided when possible". For these reasons, I wonder if there is a safer and easier way to achieve this problem.

Answer 1

You can use np.arange(50) as an indexer for both positions; it will give pairs (0, 0), (1, 1), ..., (49, 49) as the indexers for those 2 positions:

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

---

sample:

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

array([[[[  0,   0],               # changed
         [  2,   3],
         [  4,   5],
         [  6,   7],
         [  8,   9]],

        [[ 10,  11],
         [  0,   0],               # changed
         [ 14,  15],
         [ 16,  17],
         [ 18,  19]],

        [[ 20,  21],
         [ 22,  23],
         [  0,   0],               # changed
         [ 26,  27],
         [ 28,  29]],

        [[ 30,  31],
         [ 32,  33],
         [ 34,  35],
         [  0,   0],               # changed
         [ 38,  39]],

        [[ 40,  41],
         [ 42,  43],
         [ 44,  45],
         [ 46,  47],
         [  0,   0]]],             # changed


       [[[  0,   0],               # changed
         [ 52,  53],
         [ 54,  55],
         [ 56,  57],
         [ 58,  59]],

        [[ 60,  61],
         [  0,   0],               # changed
         [ 64,  65],
         [ 66,  67],
         [ 68,  69]],

        [[ 70,  71],
         [ 72,  73],
         [  0,   0],               # changed
         [ 76,  77],
         [ 78,  79]],

        [[ 80,  81],
         [ 82,  83],
         [ 84,  85],
         [  0,   0],               # changed
         [ 88,  89]],

        [[ 90,  91],
         [ 92,  93],
         [ 94,  95],
         [ 96,  97],
         [  0,   0]]],             # changed


       [[[  0,   0],               # changed
         [102, 103],
         [104, 105],
         [106, 107],
         [108, 109]],

        [[110, 111],
         [  0,   0],               # changed
         [114, 115],
         [116, 117],
         [118, 119]],

        [[120, 121],
         [122, 123],
         [  0,   0],               # changed
         [126, 127],
         [128, 129]],

        [[130, 131],
         [132, 133],
         [134, 135],
         [  0,   0],               # changed
         [138, 139]],

        [[140, 141],
         [142, 143],
         [144, 145],
         [146, 147],
         [  0,   0]]]])            # changed