numpy索引比较两个元素

时间:2017-03-25 02:09:21

标签: python numpy

我有一个1D数组,我想比较两个兄弟元素,其中元素大于阈值的条件和前一个元素小于或等于阈值 。我有一个当前的解决方案,通过数组循环比较哪个元素。这是一个简单的例子:

t = np.array([1,2,2,2,2,3,5,4,5,3,4,5,6,8,9,7,8,9,10])
threshold = 5
target_index = -1
for index in range(1, len(t)):
    if t[index] > threshold and t[index-1] <= threshold
        target_index = index
        break

我试图弄清楚是否有某种类型的花式索引解决方案或更快的方法。

1 个答案:

答案 0 :(得分:1)

我使用roll,logical_and以及where获取具有所需条件的索引。

t = np.array([1,2,2,2,2,3,5,4,5,3,4,5,6,8,9,7,8,9,10])
threshold = 5

from numpy import logical_and, roll, where

def ops_func(t, threshold):
    target_index = -1
    for index in range(1, len(t)):
        if (t[index] > threshold) and (t[index-1] <= threshold) :
            target_index = index
            break

    return target_index


def myfunc(t, threshold):
    return where(logical_and(t>threshold,roll(t,1)<=threshold))


print ops_func(t,threshold)
# lists the first index which satisfy the desired condition
print ops_func(t,threshold=3)

print myfunc(t, threshold)
# lists all the indices which satisfy the desired condition
print myfunc(t, threshold=3)

%timeit myfunc
%timeit ops_func

结果

12
6
(array([12], dtype=int64),)
(array([ 6, 10], dtype=int64),)
10000000 loops, best of 3: 23.3 ns per loop
100000000 loops, best of 3: 19.5 ns per loop