Изменение элементов в массиве NumPy

Question

Я хочу изменить все элементы в массиве A (по оси = 1) на 0, в соответствии со следующими критериями (игрушечный код):

import numpy as np

A = np.array([[1,3], [2,5], [6,2]] )
B = np.array([[1,1,0,0,0],[1,0,0,2,0],[0,0,2,2,2],[0,0,0,2,0],[6,6,0,0,0]])

for i in A:
    if i[1]<=2:
        B[B==i[0]]=0

# result

>>> B
array([[1, 1, 0, 0, 0],
       [1, 0, 0, 2, 0],
       [0, 0, 2, 2, 2],
       [0, 0, 0, 2, 0],
       [0, 0, 0, 0, 0]])

Но, к сожалению, это НЕТ 'за петли :) Спасибо!

Answer 1

В одну строку: B[np.isin(B, A[A[:, 1] <= 2][:, 0])] = 0

Объяснение:

c = A[:, 1] <= 2 # broadcast the original `if i[1]<=2:` check along axis=1
# i.e., mask A according to where the second values of the pairs are <= 2
d = c[:, 0] # index with the mask, and select the old `i[0]` values, here just `6`
e = np.isin(B, d) # mask B according to where the values are in the above
B[e] = 0 # and zero out those positions, i.e. where the old B value is 6

Answer 2

Вы можете использовать условное понимание списка, чтобы создать список первого значения в паре кортежей, где второе значение меньше или равно двум (в примере для A это последний элемент, который даетзначение 6).

Затем используйте нарезку с np.isin, чтобы найти элементы в B, содержащиеся в значениях из предыдущего условия, а затем установите эти значенияв ноль.

target_val = 2
B[np.isin(B, [a[0] for a in A if a[1] <= target_val])] = 0

>>> B
array([[1, 1, 0, 0, 0],
       [1, 0, 0, 2, 0],
       [0, 0, 2, 2, 2],
       [0, 0, 0, 2, 0],
       [0, 0, 0, 0, 0]])

В качестве альтернативы, вы также можете использовать np.where вместо нарезки.

np.where(np.isin(B, [a[0] for a in A if a[1] <= target_val]), 0, B)