Как выбрать конкретную недиагональную полосу из матрицы?

Question

Предположим, NxN квадратная матрица.Я хотел бы выбрать полосу шириной w, которая смещена на k к главной диагонали, как показано в следующем примере для N = 9, w = 4, k = -1

Порядок элементов должен быть таким же, как для np.tril_indices(N, k=k).Таким образом, в основном соответствующие индексы должны быть np.tril_indices(N, k=k) уменьшены на np.tril_indices(N, k=k-w):

i1 = list(zip(*np.tril_indices(N, k=k)))
i2 = list(zip(*np.tril_indices(N, k=k-w)))
indices = tuple(zip(*[i for i in i1 if i not in i2]))
M[indices]

или альтернативно с использованием np.tril, который, однако, сохраняет полную матрицу в качестве промежуточного результата и работает толькодля двумерных массивов:

B = np.ravel(np.tril(M, k=k) - np.tril(M, k=k-w))
B[B.nonzero()]

Мне интересно, есть ли более эффективный и / или более краткий способ добиться этого?

Answer 1

In [19]: np.tri(6,6, dtype=bool)
Out[19]: 
array([[ True, False, False, False, False, False],
       [ True,  True, False, False, False, False],
       [ True,  True,  True, False, False, False],
       [ True,  True,  True,  True, False, False],
       [ True,  True,  True,  True,  True, False],
       [ True,  True,  True,  True,  True,  True]])

Комбинирование 2 tri масок:

In [22]: np.tri(6,6, dtype=bool)&~np.tri(6,6,-3,dtype=bool)
Out[22]: 
array([[ True, False, False, False, False, False],
       [ True,  True, False, False, False, False],
       [ True,  True,  True, False, False, False],
       [False,  True,  True,  True, False, False],
       [False, False,  True,  True,  True, False],
       [False, False, False,  True,  True,  True]])

In [23]: np.where(Out[22])
Out[23]: 
(array([0, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5]),
 array([0, 0, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5]))

np.tri по существу:

In [29]: np.greater_equal.outer(np.arange(4),np.arange(4))
Out[29]: 
array([[ True, False, False, False],
       [ True,  True, False, False],
       [ True,  True,  True, False],
       [ True,  True,  True,  True]])
In [30]: np.greater_equal(np.arange(4)[:,None],np.arange(4))
Out[30]: 
array([[ True, False, False, False],
       [ True,  True, False, False],
       [ True,  True,  True, False],
       [ True,  True,  True,  True]])

np.tril только применимо where к этому.

In [37]: np.greater_equal(np.arange(4)[:,None], np.arange(4)) & 
         np.less_equal(np.arange(-1,3)[:,None], np.arange(4))
Out[37]: 
array([[ True, False, False, False],
       [ True,  True, False, False],
       [False,  True,  True, False],
       [False, False,  True,  True]])