Как найти максимум столбца в 3D-массиве на основе условия в другом столбце?

Question

У меня есть массив с 3 столбцами, и мне нужно найти максимум первого столбца, где третий столбец меньше некоторого предопределенного значения.

Ввод, напр .:

a = [[1 2 1], [2 3 3], [7 10 5], [5 1 2], [2 3 4]]

Желаемый вывод

max (a [:, 0], где a [:, 2] <5) = [5 1 2] </p>

Как это сделать? Желательно как можно быстрее.

Answer 1

Вы можете использовать np.argmax:

subset = a[a[:, 2] < 5]
result = subset[subset[:, 0].argmax()]

print(result)

Вывод:

[5 1 2]

Получает подмножество исходного массива на основе заданного условия, находит индекс строкис самым высоким значением в первом столбце с нарезкой и argmax, и использует это значение для индексации подмножества.

Answer 2

Вы можете сделать с argmax

a[(a[:,0] * (a[:,2] < 5).astype(int)).argmax(),:]
Out[946]: array([5, 1, 2])