Numpy equivalent of list.index

前端 未结 6 2250
粉色の甜心
粉色の甜心 2021-02-12 11:48

In a low-level function that is called many times, I need to do the equivalent of python\'s list.index, but with a numpy array. The function needs to return when it finds the f

6条回答
  •  北荒
    北荒 (楼主)
    2021-02-12 12:24

    NumPy's searchsorted is very similar to lists's index, except that it requires a sorted array and behaves more numerically. The big differences are that you don't need to have an exact match, and you can search starting from either the left or right sides. See the following examples to get an idea how it works:

    import numpy as np
    a = np.array([10, 20, 30])
    
    a.searchsorted(-99) == a.searchsorted(0) == a.searchsorted(10)
    # returns index 0 for value 10
    
    a.searchsorted(20.1) == a.searchsorted(29.9) == a.searchsorted(30)
    # returns index 2 for value 30
    
    a.searchsorted(30.1) == a.searchsorted(99) == a.searchsorted(np.nan)
    # returns index 3 for undefined value
    

    With the last case, where an index of 3 is returned, you can handle this as you like. I gather from the name and intention of the function that it stops after finding the first suitable index.

提交回复
热议问题