кусочки, использующие индекс конца столбца из целочисленного массива

Question

Если у меня есть массив и я применяю суммирование

arr = np.array([[1.,1.,2.],[2.,3.,4.],[4.,5.,6]])
np.sum(arr,axis=1)

Я получаю сумму по трем строкам ([4., 9., 15.])

Моя сложностьэтот arr содержит данные, которые могут быть плохими после определенного индекса столбца.У меня есть целочисленный массив, который говорит мне, сколько «хороших» значений у меня есть в каждой строке, и я хочу суммировать / усреднять по хорошим значениям.Скажите:

ngoodcols=np.array([0,1,2])
np.sum(arr[:,0:ngoodcols],axis=1)  # not legit but this is the idea

Понятно, как сделать это в цикле, но есть ли способ суммировать только это число, получая [0., 2., 9.] Без обращения к циклу?Эквивалентно, я мог бы использовать nansum, если бы знал, как установить элементы в индексах столбцов выше, чем b, равный np.nan, но это почти эквивалентная проблема с точки зрения нарезки.

Answer 1

Одной из возможностей является использование масочных массивов :

import numpy as np

arr = np.array([[1., 1., 2.], [2., 3., 4.], [4., 5., 6]])
ngoodcols = np.array([0, 1, 2])
mask = ngoodcols[:, np.newaxis] <= np.arange(arr.shape[1])
arr_masked = np.ma.masked_array(arr, mask)
print(arr_masked)
# [[-- -- --]
#  [2.0 -- --]
#  [4.0 5.0 --]]
print(arr_masked.sum(1))
# [-- 2.0 9.0]

Обратите внимание, что здесь, когда нет хороших значений, в результате вы получаете «пропущенное» значение, которое может или не может быть полезным для вас. Кроме того, замаскированный массив также позволяет вам легко выполнять другие операции, которые применяются только к допустимым значениям (среднее и т. Д.).

Другой простой вариант - просто умножить на маску:

import numpy as np

arr = np.array([[1., 1., 2.], [2., 3., 4.], [4., 5., 6]])
ngoodcols = np.array([0, 1, 2])
mask = ngoodcols[:, np.newaxis] <= np.arange(arr.shape[1])
print((arr * ~mask).sum(1))
# [0. 2. 9.]

Здесь, когда нет хороших значений, вы просто получаете ноль.

Answer 2

Вот один из способов использования логического индексирования. Это устанавливает элементы в индексах столбцов выше, чем в ngoodcols равными np.nan и использует np.nansum:

import numpy as np

arr = np.array([[1.,1.,2.],[2.,3.,4.],[4.,5.,6]])
ngoodcols = np.array([0,1,2])

arr[np.asarray(ngoodcols)[:,None] <= np.arange(arr.shape[1])] = np.nan

print(np.nansum(arr, axis=1))
# [ 0.  2.  9.]