获取Numpy数组的索引

时间:2018-05-11 20:27:16

标签: python arrays numpy search

我有一个numpy数组:

arr = [0.23, 2.32, 4.04, 5.02, 6.84, 10.12, 10.34, 11.93,12.44]

我想得到我输入的最接近整数的索引。例如,如果我输入10然后我应该回到索引5(10.12)或如果我输入12我应该回到索引7(11.93)。

1 个答案:

答案 0 :(得分:2)

如果您的列表未排序,则需要使用abs + argmin的线性时间解决方案:

>>> np.abs(np.array(arr) - 12).argmin()
7

但是,如果您的列表已排序(升序或降序),您可以使用二进制搜索进行亚线性时间解决方案(非常快):

# https://ideone.com/aKEpI2 — improved by @user2357112
def binary_search(arr, val):
    # val must be in the closed interval between arr[i-1] and arr[i],
    # unless one of i-1 or i is beyond the bounds of the array.
    i = np.searchsorted(arr, val)

    if i == 0:
        # Smaller than the smallest element
        return i
    elif i == len(arr):
        # Bigger than the biggest element
        return i - 1
    elif val - arr[i - 1] <= arr[i] - val:
        # At least as close to arr[i - 1] as arr[i]
        return i - 1

    # Closer to arr[i] than arr[i - 1]
    return i

cases = [10, 12, 100, 10.12]   # 5, 7, 8, 5
print(*[binary_search(arr, c) for c in cases], sep=',')