提取np.diagonal的位置

Question

我有一个二维的numpy数组，我想用np.diagonal来提取对角线上的元素。更准确地说，我不是想要实际的值，而是想要那些np.diagonal会提取的值的位置。

size = 10
value = 1
offset = 1

array = np.random.randint(0, 10, size=(size, size))

diagonal = np.diag(np.full(size,value))+np.diag(np.full(size-1, value),offset)+np.diag(np.full(size-1, value), -offset)   
diagonal = np.where(diagonal==0, np.nan, diagonal)

如果偏移量是1，这个方法效果很好。但如果偏移量不是1，数组的大小就会不同，这样就会出现错误：

ValueError: operands could not be broadcast together with shapes (10,10) (11,11) 

我该怎么解决这个问题呢？最终得到的diagonal数组应该和原来的array数组大小相同。

Answer 1

你可以使用 np.indices 来解决这个问题：

import numpy as np
size = 10
value = 1
offset = 2

x, y = np.indices((size, size))
np.where((abs(x - y) == offset) | (x == y), value, np.nan)

array([[ 1., nan,  1., nan, nan, nan, nan, nan, nan, nan],
       [nan,  1., nan,  1., nan, nan, nan, nan, nan, nan],
       [ 1., nan,  1., nan,  1., nan, nan, nan, nan, nan],
       [nan,  1., nan,  1., nan,  1., nan, nan, nan, nan],
       [nan, nan,  1., nan,  1., nan,  1., nan, nan, nan],
       [nan, nan, nan,  1., nan,  1., nan,  1., nan, nan],
       [nan, nan, nan, nan,  1., nan,  1., nan,  1., nan],
       [nan, nan, nan, nan, nan,  1., nan,  1., nan,  1.],
       [nan, nan, nan, nan, nan, nan,  1., nan,  1., nan],
       [nan, nan, nan, nan, nan, nan, nan,  1., nan,  1.]])

Answer 2

最简单的解决办法似乎是从你用来构建对角线的数组长度中减去偏移量。

diagonal = np.diag(np.full(size,value))+np.diag(np.full(size-offset, value),offset)+np.diag(np.full(size-offset, value), -offset)   
diagonal = np.where(diagonal==0, np.nan, diagonal)

偏移量为2时的输出：

array([[ 1., nan,  1., nan, nan, nan, nan, nan, nan, nan],
       [nan,  1., nan,  1., nan, nan, nan, nan, nan, nan],
       [ 1., nan,  1., nan,  1., nan, nan, nan, nan, nan],
       [nan,  1., nan,  1., nan,  1., nan, nan, nan, nan],
       [nan, nan,  1., nan,  1., nan,  1., nan, nan, nan],
       [nan, nan, nan,  1., nan,  1., nan,  1., nan, nan],
       [nan, nan, nan, nan,  1., nan,  1., nan,  1., nan],
       [nan, nan, nan, nan, nan,  1., nan,  1., nan,  1.],
       [nan, nan, nan, nan, nan, nan,  1., nan,  1., nan],
       [nan, nan, nan, nan, nan, nan, nan,  1., nan,  1.]])

如果你真的需要位置的话，可以把

diagonal = np.where(diagonal==0, np.nan, diagonal)

改成

i, j = np.where(diagonal)

编辑：

这里有一个替代方案，可以创建主对角线和偏移对角线之间的所有对角线。

size = 10
offset = 1
A = np.ones((size, size))
A[np.tril_indices(size, k=-(offset + 1))] = 0
A[np.triu_indices(size, k=(offset + 1))] = 0
A = np.where(A, 1, np.nan)