Найти локальный минимум в 2D np.array в Python

Question

У меня есть этот массив из двух столбцов:

A  | 1 
A  | 2
A  | 3
B  | 4
B  | 5
B  | 6

, где A, B - константы. Я хочу найти минимальное значение каждого параметра A и B, поэтому результатом этой операции будет другой двумерный массив, подобный этому:

A | 1
B | 4

Мне удастся найти mimium, когда только один присутствует константа A:

MIN = np.where(arr == np.amin(arr[:,1]))
output = arr[MIN[0],:]
>> output = A | 1

Но я не могу автоматизировать ее, когда массив более сложен, чем этот. Спасибо за помощь.

Answer 1

Вы можете использовать np.unique для извлечения уникальных значений из первого столбца, а затем использовать логическую маску вдоль первой оси:

np.stack([(x, a[a[:, 0] == x, 1].min()) for x in np.unique(a[:, 0])])

В качестве альтернативы, если у вас есть доступ к pandas, вы может создать DataFrame и использовать grouby в первом столбце:

import pandas as pd

result = pd.DataFrame(a).groupby(0).min().values

---

Код, используемый, например, для установки a:

import numpy as np

a = np.array([[1, 1, 1, 2, 2, 2],
              [1, 2, 3, 4, 5, 6]]).T

Answer 2

То, что вы ищете, это необязательное ключевое слово axis в функции np.min. Это позволяет вычислять минимум массива по столбцам. Использование np.min также лучше, чем использование np.amin, поскольку оно позволяет выполнить на один шаг меньше до результата (у вас сразу же есть минимум вместо его индекса в таблице)

Попробуйте это:

import numpy as np

a = np.array([[1,2,3], [4,5,6]])
np.min(a, axis=1)
# Result : array([1, 4])