Numpy 数组条件检查是否存在值为当前值 +1 的元素

Question

我想将类似这样的 if 语句构建到 numpy 条件中（用于 argwhere 或类似条件）。

for i, element in enumerate(array):
    if array[i]+1 == array[i+1]:
        return True

这是可能的还是遍历数组的唯一方法？

Answer 1

您可以使用 np.roll 移动数组的函数

import numpy as np
array+1 == np.roll(array, shift=-1)

这将 return 一个布尔值数组。

Answer 2

您也可以使用 np.equal() 功能。

import numpy as np

a = [1,2,4,8,12,13]
x = np.array(a)
res = np.equal((x+1)[:-1], x[1:])

print(res)

输出：

[ True False False False  True]

注：

如果你需要速度，通常情况下，当我们使用 numpy 时，值得一提的是，这种方法比 np.roll() 更快，FBruzzesi 在下面提出了什么，这是一个优雅的解决方案无论如何：

import timeit

code1 = """
import numpy as np
a = [1,2,4,8,12,13]
x = np.array(a)
np.equal((x+1)[:-1], x[1:])
"""
elapsed_time1 = timeit.timeit(code1, number=10000)/100
print(elapsed_time1)

code2 = """
import numpy as np
a = [1,2,4,8,12,13]
x = np.array(a)
x+1 == np.roll(x, shift=-1)
"""

elapsed_time2 = timeit.timeit(code2, number=10000)/100
print(elapsed_time2)

输出：

0.00044608700000026147
0.0022752689999970244

Numpy 数组条件检查是否存在值为当前值 +1 的元素

Numpy array condition to check if there is element with value +1 of current value

python

loops

numpy

vectorization