numpy 2d: как получить индекс максимального элемента в первом столбце только для разрешенного значения во втором столбце

#python #arrays #numpy #max #masked-array

#python #массивы #numpy #max #masked-array

Вопрос:

Помогите найти высокопроизводительный способ решения проблемы: у меня есть результат после neural-network (answers_weight), категория для ответов (тот же len) и разрешенные категории для текущего запроса:

 answers_weight = np.asarray([0.9, 3.8, 3, 0.6, 0.7, 0.99]) # ~3kk items
answers_category = [1, 2, 1, 5, 3, 1] # same size as answers_weight: ~3kk items
categories_allowed1 = [1, 5, 8]
res = np.stack((answers_weight, answers_category), axis=1)
  

Мне нужно знать индекс (в массиве answers_weight) максимального элемента, но пропустить недопустимые категории (2,3).

В финале индекс должен быть = 2 («3.0», потому что «3.8» должно быть пропущено как недопустимое по категории)

Ответ №1:

Самым простым способом было бы использовать numpy masked_arrays для маскирования ваших весов в соответствии с allowed_categories, а затем найти argmax :

 np.ma.masked_where(~np.isin(answers_category,categories_allowed1),answers_weight).argmax()
#2
  

Другой способ сделать это с помощью масок (этот предполагает уникальный максимальный вес):

 mask = np.isin(answers_category, categories_allowed1)
np.argwhere(answers_weight==answers_weight[mask].max())[0,0]
#2
  

Ответ №2:

Я также решил эту проблему с помощью маски

 inds = np.arange(res.shape[0])
# a mask is an array [False  True False False  True False]
mask = np.all(res[:,1][:,None] != categories_allowed1,axis=1)

allowed_inds = inds[mask]
# max_ind is not yet the real answer because the not allowed values are not taken into account
max_ind = np.argmax(res[:,0][mask])
real_ind = allowed_inds[max_ind]