在标签值之前获取 n 个最高项目的索引

Question

我有一个温度列表，我需要在所需温度之前和之后找到 n 个值，最好在单独的列表中。我列表中的值不一定是唯一的，但我需要原始列表的索引。我需要这些索引来查找其他列表中的其他参数。

例子：

TestArray = np.array([12,42,19,32,41,14,17,23,12,18,32,19])
Value = 20
n = 2
TestArray = np.append(TestArray, Value) 
Sort = np.argsort(TestArray)
Index = np.where(Sort == (len(TestArray)-1))[0][0]
Lower = Sort[Index-n:Index]
Upper = Sort[Index+1:Index+n+1]
print(Upper, TestArray[Upper])
print(Lower, TestArray[Lower])

我的代码给出了想要的输出，但它看起来真的很乱，我想知道是否有更好的方法。

说明： 我将想要的值附加到末尾，所以我知道它的索引。然后我使用 argsort 从低到高查找索引，然后使用 np.where 定位我想要的值在哪里。然后使用索引查找上限值和下限值。

Answer 1

你可以过滤数组然后使用np.partition:

mask = TestArray < Value
Lower, Upper = -np.partition(-TestArray[mask], 2)[:2], np.partition(TestArray[~mask], 2)[:2]

Lower
#array([19, 19])

Upper
#array([23, 32])

要取回索引：

TestArray = np.array([12,42,19,32,41,14,17,23,12,18,32,19])    ​
​
mask = TestArray < Value
arrInd = np.column_stack((np.arange(len(TestArray)), TestArray))
Lower, Upper = arrInd[mask,:], arrInd[~mask,:]
LowerInd, UpperInd = np.argpartition(-Lower[:, 1], 2)[:2], np.argpartition(Upper[:,1], 2)[:2]
​
print(Lower[LowerInd])
​#[[ 2 19]
# [11 19]]

print(Upper[UpperInd])
#[[ 7 23]
# [ 3 32]]

Answer 2

我认为np.searchsorted 可能对你有用。

这是来自documentation的示例

>>> np.searchsorted([1,2,3,4,5], 3)
2

Answer 3

numpy 解决方案可能是最好的方法。作为比较，这里有一个使用heapq.nsmallest() 的有效解决方案：

>>> from heapq import nsmallest
>>> data = [12,42,19,32,41,14,17,23,12,18,32,19]
>>> nsmallest(2, data, key=lambda x: (x-20 if x >= 20 else float('inf')))
[23, 32]
>>> nsmallest(2, data, key=lambda x: (20-x if x <= 20 else float('inf')))
[19, 19]