Массив Numpy: как по строкам проверять правильность первого значения X?

Question

Описание проблемы

Рассмотрим следующие два примера массивов:

arr = np.array([
    [5.0, 2.0, 1.0, np.nan, np.nan],
    [9.0, np.nan, np.nan, np.nan, 2.0],
    [4.0, 7.0, 4.0, np.nan, np.nan],
    [8.0, np.nan, np.nan, np.nan, np.nan],
    [np.nan, np.nan, np.nan, np.nan, np.nan],
    [np.nan, np.nan, np.nan, np.nan, 6.0]
])

amounts = np.array([
    3,
    1,
    2,
    3,
    0,
    5
])

Для каждой строки в массиве arr Я хотел бы проверить, не являются ли первые X записей в строке неNaN, но все остальные записи NaN.Это количество X различно для каждой строки и задается массивом amounts.

Таким образом, мой ожидаемый результат будет следующим логическим массивом:

array([ True, False, False, False,  True, False])

Пробовал до сих пор

Мне удалось придумать следующий рабочий код:

result = []
for (row, amount) in zip(arr, amounts):
    if (~np.isnan(row)[:amount]).all() and np.isnan(row)[amount:].all():
        result.append(True)
    else:
        result.append(False)

result = np.array(result)
print(result)

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

Кто-нибудь сможет помочь найти полностью векторизованное решение этой проблемы?

Answer 1

a = np.array([[5.0, 2.0, 1.0, np.nan, np.nan],
              [9.0, np.nan, np.nan, np.nan, 2.0],
              [4.0, 7.0, 4.0, np.nan, np.nan],
              [8.0, np.nan, np.nan, np.nan, np.nan],
              [np.nan, np.nan, np.nan, np.nan, np.nan],
              [np.nan, np.nan, np.nan, np.nan, 6.0]])

b = np.array([3,1,2,3,0,5])

c = np.logical_not(np.isnan(a))
firstn = b == c.argmin(axis=1)
no_extras = b == c.sum(axis=1)
result = np.logical_and(firstn,no_extras)

Создание логического массива значений, отличных от NaN.

Чтобы обеспечить соответствие первых значений n критериям;Используйте numpy.argmin(), чтобы найти первое NaN - сравните это с массивом counts.

Чтобы убедиться, что нет никаких значений, отличных от NaN после NaNs начать ;Суммируйте все True s по строкам в логическом массиве, сравните это с массивом counts.

and оба из этих результатов.

Answer 2

Вы можете попробовать что-то вроде этого:

# Values are column numbers
grid = np.tile(np.arange(arr.shape[1]), (arr.shape[0], 1))

# Mask
mask = grid < amounts.reshape((-1, 1))

# Comparison
np.all(~np.isnan(arr) == mask, axis=1)