Numpy: как установить столбец [index:] = значение, для каждого столбца, для массива различных индексов?

Question

Предположим, у меня есть массив:

my_array = np.random.normal(size=(5,3))
print(my_array)

[[ 0.45110035 -1.08385534  1.2126054 ]
 [ 1.51280316  0.4308235  -0.31839059]
 [-0.00348102 -0.50814392  0.00734745]
 [-0.63701191  0.95413945 -1.40480595]
 [-1.66723431 -0.52822503 -1.14282036]]

И список индексов той же длины, что и количество столбцов:

my_indices = np.array([3, 1, 2])

Для каждого столбца в my_array,Я хочу получить соответствующий индекс из my_indices и установить для этого элемента и всех последующих элементов этого столбца какое-то новое значение.Другими словами, результат будет выглядеть следующим образом:

my_new_array = np.something(my_array, my_indices, 23)
print(my_new_array)

[[ 0.45110035 -1.08385534  1.2126054 ]
 [ 1.51280316  23.0       -0.31839059]
 [-0.00348102  23.0        23.0      ]
 [ 23.0        23.0        23.0      ]
 [ 23.0        23.0        23.0      ]]
# i.e.:
# col[3:]=23  col[1:]=23  col[2:]=23

Какой самый простой способ добиться этого?Я знаю, что могу сделать это с помощью простого цикла for, но это плохо для numpy.Моя интуиция заключается в том, что, возможно, есть какой-то умный быстрый способ сделать это с помощью трансляции и векторизации, но я не могу понять это.

Answer 1

Установить или выбрать один элемент из каждого столбца легко:

In [10]: arr = np.ones((5,3),int)
In [11]: b = np.array([3,1,2])
In [12]: arr[b, np.arange(3)] = 0
In [13]: arr
Out[13]: 
array([[1, 1, 1],
       [1, 0, 1],
       [1, 1, 0],
       [0, 1, 1],
       [1, 1, 1]])

Но установить фрагмент для каждого столбца сложнее.Самое простое - просто итерировать по столбцам.

In [14]: for i,j in enumerate(b):
    ...:     arr[j:, i] = 0
    ...:     
In [15]: arr
Out[15]: 
array([[1, 1, 1],
       [1, 0, 1],
       [1, 0, 0],
       [0, 0, 0],
       [0, 0, 0]])

Другой способ заключается в настройке маски с использованием сравнения b с np.arange(5).Это быстрее, но не так интуитивно понятно.Я должен переосмыслить подход каждый раз.

In [16]: mask = np.arange(5)[:,None]>=b
In [17]: mask
Out[17]: 
array([[False, False, False],
       [False,  True, False],
       [False,  True,  True],
       [ True,  True,  True],
       [ True,  True,  True]])
In [18]: arr[mask] = 2
In [19]: arr
Out[19]: 
array([[1, 1, 1],
       [1, 2, 1],
       [1, 2, 2],
       [2, 2, 2],
       [2, 2, 2]])