在tage值之前获取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]