Фильтр numpy массива в соответствии с предикатом позиции

Question

Предположим, у меня есть двумерный массив numpy с заданным shape, и я хотел бы получить представление о значениях, которые удовлетворяют предикату, основанному на позиции значения. То есть, если x и y соответствуют индексу столбца и строки, а предикат x>y, функция должна возвращать только значения массива, для которых индекс столбца больше, чем индекс строки.

Простой способ сделать это двойной цикл, но я бы хотел, возможно, более быстрый (векторизованный, возможно?) Подход?

Есть ли лучший способ?

Answer 1

В общем, вы можете сделать это путем построения сетки открытой сетки, соответствующие индексы строк/столбцов, применить предикат, чтобы получить булеву маску, то индекс в массив, используя эту маску:

A = np.zeros((10,20)) 
y, x = np.ogrid[:A.shape[0], :A.shape[1]] 
mask = x > y 
A[mask] = 1 

Ваш конкретный пример - это верхний треугольник - вы можете получить его копию с помощью np.triu, или вы можете получить соответствующие индексы строк/столбцов, используя np.triu_indices.