Замена элементов двумерного массива под индексным массивом

Question

Я искал предыдущие вопросы и не нашел того, который вполне подходит для этого.

Я устанавливаю значение для всех ячеек j -го столбца A на 0, чьи значения строки меньше m [j]:

import numpy as np
n = 8
A = np.ones([n,n])
m = np.array([1,1,1,2,2,2,3,4])
for j in range(n):
    for i in range(m[j]):
        A[i,j] = 0

Как мне переписать этот скрипт без циклов for?

Answer 1

Вы можете использовать транслируемое сравнение NumPy -

>>> A[np.arange(n)[:, None] < m] = 0
>>> 
array([[0., 0., 0., 0., 0., 0., 0., 0.],
       [1., 1., 1., 0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1., 1., 0., 0.],
       [1., 1., 1., 1., 1., 1., 1., 0.],
       [1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1.]])

Здесь вызов [:, None] увеличивает форму np.arange(n), так что сравнение < транслируется по каждому элементу m для каждого элемента в диапазоне. При этом генерируется логическая маска той же формы, что и A, которая затем используется для установки значений на 0.

Примечание. Если гарантировано, что A будет массивом из них, я бы порекомендовал решение Дивакара , которое очень похоже на это.

Answer 2

Использование внешнего широковещательного сравнения с ранжированным массивом длины n для прямого получения A -

A = (m <=  np.arange(n)[:,None]).astype(float)

Чтобы повысить производительность, мы могли бы использовать преобразование типов для более точныхкак .astype(np.int) или даже .astype(np.uint8).Если у вас все в порядке с логическим массивом, пропустите преобразование типов вообще.

Для int dtypes, мы также можем просто view в него, что экономит нам немного памяти -

(m <=  np.arange(n)[:,None]).view(np.int8)