Эффективно назначьте ноль нескольким столбцам в каждой строке двумерного массива

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

Например, здесь мы имеем матрицу A размером 4x5 с обрезанными столбцами [1,3,2,4]. Что я хочу сделать:

import numpy as np
np.random.seed(1)
A = np.random.rand(4, 5)
cutoff = np.array([1,3,2,4])
A[0, cutoff[0]:] = 0
A[1, cutoff[1]:] = 0
A[2, cutoff[2]:] = 0
A[3, cutoff[3]:] = 0

Я мог бы сделать это с помощью np.repeat ing индексов строк и столбцов, но моя матрица настолько велика, что я не могу этого сделать. Есть ли эффективный способ сделать это?

1 ответ

Решение

Использование broadcasting создать полную маску и назначить -

A[cutoff[:,None] <= np.arange(A.shape[1])] = 0

Альтернативно с builtin внешний метод -

A[np.less_equal.outer(cutoff, range(A.shape[1]))] = 0
Другие вопросы по тегам