Как выбрать столбцы из матрицы в NumPy со значениями из другого вектора

Question

Как выбрать столбцы на основе вектора, примененного к матрице? Расширение: Как применить маску из массива к другой матрице в numpy

ar2 = np.arange(1,26)[::-1].reshape([5,5]).T
ar3 = np.array([1,2,3,1,2])
print (ar2, '\n\n',  ar3)

results in:

[[25 20 15 10  5]
 [24 19 14  9  4]
 [23 18 13  8  3]
 [22 17 12  7  2]
 [21 16 11  6  1]] 

 [1 2 3 1 2]

Что я ищу после:

1  [[ 25  np.nan np.nan np.nan np.nan]
2   [ 24   19    np.nan np.nan np.nan]
3   [ 23   18      13   np.nan np.nan]
1   [ 22  np.nan np.nan np.nan np.nan]
2   [ 21   16    np.nan np.nan np.nan ]] 

Answer 1

Мы можем использовать broadcasting, чтобы создать маску со сравнением массива на расстоянии с ar3 для назначения в эти места, а затем назначить NaNs. Поскольку входные данные представляют собой массив int, нам нужно сделать плавающую копию ar2 и затем присвоить, например, -

out = ar2.astype(float, copy=True) # convert to float as NaNs are to be assigned
mask = ar3[:,None] <= np.arange(ar2.shape[1])
out[mask] = np.nan

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

---

Немного больше объяснений по созданию mask -

In [38]: ar3
Out[38]: array([1, 2, 3, 1, 2])

In [39]: ar3[:,None] <= np.arange(ar2.shape[1])
Out[39]: 
array([[False,  True,  True,  True,  True],
       [False, False,  True,  True,  True],
       [False, False, False,  True,  True],
       [False,  True,  True,  True,  True],
       [False, False,  True,  True,  True]])

Сравнение каждого элемента ar3 с range(5) с этим outer сравнением дает нам каждый ряд маски. Если мы посмотрим внимательнее, это все False до соответствующего значения индекса (значение ar3) и True после этого. Нам нужны эти True мест для назначения NaNs, и, следовательно, эта маска напрямую помогает нам в определении NaN во всем выходном массиве.