python之将具有匹配索引的数组元素清零

我想将 n 维数组中位于具有两个或多个匹配索引的位置的所有元素“归零”。在二维中，这实际上是 np.fill_diagonal() 但当考虑第三个维度时，该函数就变得不够用了。

下面是我想做的暴力破解版本。有什么方法可以清理它并使其在 n 个维度上工作吗？

x = np.ones([3,3,3]) 
 
x[:,0,0] = 0 
x[0,:,0] = 0 
x[0,0,:] = 0 
 
x[:,1,1] = 0 
x[1,:,1] = 0 
x[1,1,:] = 0 
 
x[:,2,2] = 0 
x[2,:,2] = 0 
x[2,2,:] = 0 
 
print(x) 

请您参考如下方法:

一种方法是np.einsum:

>>> a = np.ones((4,4,4), int) 
>>> for n in range(3): 
...     np.einsum(f"{'iijii'[n:n+3]}->ij", a)[...] = 0 
...  
>>> a 
array([[[0, 0, 0, 0], 
        [0, 0, 1, 1], 
        [0, 1, 0, 1], 
        [0, 1, 1, 0]], 
 
       [[0, 0, 1, 1], 
        [0, 0, 0, 0], 
        [1, 0, 0, 1], 
        [1, 0, 1, 0]], 
 
       [[0, 1, 0, 1], 
        [1, 0, 0, 1], 
        [0, 0, 0, 0], 
        [1, 1, 0, 0]], 
 
       [[0, 1, 1, 0], 
        [1, 0, 1, 0], 
        [1, 1, 0, 0], 
        [0, 0, 0, 0]]]) 

一般(ND)情况:

>>> from string import ascii_lowercase 
>>> from itertools import combinations 
>>>  
>>> a = np.ones((4,4,4,4), int) 
>>> n = a.ndim 
>>> ltrs = ascii_lowercase[:n-2] 
>>> for I in combinations(range(n), 2): 
...     li = iter(ltrs) 
...     np.einsum(''.join('z' if k in I else next(li) for k in range(n)) + '->z' + ltrs, a)[...] = 0 
...  
>>> a 
array([[[[0, 0, 0, 0], 
         [0, 0, 0, 0], 
         [0, 0, 0, 0], 
         [0, 0, 0, 0]], 
 
        [[0, 0, 0, 0], 
         [0, 0, 0, 0], 
         [0, 0, 0, 1], 
         [0, 0, 1, 0]], 
 
        [[0, 0, 0, 0], 
         [0, 0, 0, 1], 
         [0, 0, 0, 0], 
         [0, 1, 0, 0]], 
 
<snip>