基于2D numpy数组的第二列中的最大值来查找1D numpy数组

Question

我目前正在从2D数组中删除一些基于其中一列值的1D数组。第一列可能有不同的重复值，我想根据第二列的最大值保留每个重复值中的一个(这只是一个例子，2d数组可能更大)这是我尝试的

import numpy as np
arr = np.array([[ 36.06, 209.14],
[ 36.06, 214.55],
[ 36.06, 215.91],
[ 36.06, 225.29],
[ 41.11, 186.76],
[ 41.11, 191.79],
[ 41.11, 197.21],
[ 41.11, 197.33],
[ 41.11, 201.19],
[ 41.11, 206.15],
[ 50.25, 165.51],
[ 50.25, 174.32],
[ 59.03, 148.79]])     
biggest = 0
aux = []
for i in range(arr.shape[0]-1):
j = i+1
if (arr[i][0] == arr[j][0]):
if (arr[i][1] < arr[j][1] and arr[j][1] > biggest):
biggest = j
if (arr[i][0] != arr[j][0]):
aux.append(arr[biggest])
print(np.array(aux))
#Output = [[ 36.06 225.29]
#          [ 41.11 206.15]
#          [ 50.25 174.32]]

正如你所看到的，我几乎得到了预期的结果，我的预期结果应该是这样的。。。

Output = [[ 36.06 225.29]
[ 41.11 206.15]
[ 50.25 174.32]
[ 59.03 148.79]]

问题是我缺少最后一个数组，也许有一种更简单的方法可以使用我缺少的numpy内置函数。提前谢谢！

Answer 1

没有理由重新发明轮子。只用熊猫。

import pandas as pd
pd.DataFrame(arr).groupby(0, as_index=False).max().to_numpy()
>> array([[ 36.06, 225.29],
[ 41.11, 206.15],
[ 50.25, 174.32],
[ 59.03, 148.79]])

备选方案

输入似乎在两列中都进行了排序，这意味着每个键的最高值总是最后一个。如果是这样的话，或者如果可以通过排序来完成，那么普通的numpy版本也是可能的。

# if not already sorted, sort as described above
sorted_array = arr[np.lexsort((arr[:, 1], arr[:, 0]))]
# find the last value per key
keys = sorted_array[:, 0]
ends = np.append(keys[1:] != keys[:-1], True)
# extract rows
return sorted_array[ends]

如果我们包括排序成本，这比panda版本的计算复杂度更高(假设panda版本使用哈希表；尚未检查)数据的形状和实现的质量可能会改变实际运行时。

Answer 2

一种方法是在第一列应用np.unique来查找该列中的唯一值(注意，默认情况下，np.unique将在排序方案中获得唯一值，这在您的示例中起作用)，然后检查第二列中每个唯一值的最大值索引，并将其附加到列表中：

aux = []
for i in np.unique(arr[:, 0]):
arr_ = arr[arr[:, 0] == i]
aux.append(arr_[arr_[:, 1].argmax()])

或者使用数组代替列表附加：

uniques_ = np.unique(arr[:, 0])
# [36.06 41.11 50.25 59.03]
result = np.empty((uniques_.shape[0], arr.shape[1]))
for i, j in enumerate(uniques_):
arr_ = arr[arr[:, 0] == j]
result[i] = arr_[arr_[:, 1].argmax()]
# result
# [[ 36.06 225.29]
#  [ 41.11 206.15]
#  [ 50.25 174.32]
#  [ 59.03 148.79]]

使用np.unique保留第一列的排序，如果我们有：

arr = np.array([[ 41.11, 186.76],
[ 41.11, 191.79],
[ 41.11, 197.21],
[ 41.11, 197.33],
[ 41.11, 201.19],
[ 41.11, 206.15],
[ 36.06, 209.14],
[ 36.06, 214.55],
[ 36.06, 215.91],
[ 36.06, 225.29],
[ 50.25, 165.51],
[ 50.25, 174.32],
[ 59.03, 148.79]])
_, idx = np.unique(arr[:, 0], return_index=True)
uniques_ = arr[:, 0][np.sort(idx)]
result = np.empty((uniques_.shape[0], arr.shape[1]))
for i, j in enumerate(uniques_):
arr_ = arr[arr[:, 0] == j]
result[i] = arr_[arr_[:, 1].argmax()]
# result
# [[ 41.11 206.15]
#  [ 36.06 225.29]
#  [ 50.25 174.32]
#  [ 59.03 148.79]]