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）情况：

>>来自字符串导入ascii\u小写
>>>从itertools导入组合
>>> 
>>>a=np.one（（4,4,4,4），int）
>>>n=a.ndim
>>>ltrs=ascii_小写[：n-2]
>>>对于组合中的I（范围（n），2）：
...     li=国际热核实验堆（ltrs）
...     np.einsum（“”.join（'z'如果k在I中，则下一个（li）表示k在范围（n））+'->z'+ltrs，a）[……]=0
... 
>>>a
数组（[[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]],

对于您提到的案例，似乎是一个解决方案。非常巧妙地使用f-Strings Hanks a ton Paul！但是，当维度小于4维时，一般解决方案似乎会中断。有什么想法吗？@drew_很好？是的，我在For循环中硬编码了

4

。那必须是

n

。当然！谢谢您的快速回复和回复这是一个奇妙的解决方案！

>>> 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>