在 Python 中优化 for 循环以更快地工作

Question

我正在努力优化 Python 代码。目标是获取一个整数列表并计算并输出列表中有多少对。一对被认为是 2 个数字，差异为 K（在这种情况下为 2） 例如：

k = 2
list = [1, 5, 3, 4, 2]

这里的配对将是(1,3), (5,3), (2,4)
答案是：3

我想提高代码效率，当前版本需要 8 秒或更长时间。

cProfile 告诉我for number in sorted_array: 是唯一需要花费所有时间的行。但我似乎无法弄清楚如何优化for 循环。

有没有人有任何经验或建议？非常感谢。

代码：

#generate random numbers
import bisect
import random
n_integers = random.sample(xrange(1, 29999), 29998)
####cProfile
import cProfile
pr = cProfile.Profile()
pr.enable()

#the difference between numbers we are looking for
k = 2
sorted_array = []
pairs_counter = 0

#insert N integers in array in sorted fashion and typecast
for number in n_integers:
    bisect.insort_left(sorted_array, number)

#iterate over the array and calculate (number + K)
for number in sorted_array:
    the_pair = number + k
    #check if the number+K is in the array
    if the_pair in sorted_array:
        pairs_counter += 1

print pairs_counter

#Close cProfile
pr.disable()
pr.print_stats(sort = 'time')

c个人资料：

30075 function calls in 7.995 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    7.834    7.834    7.834    7.834 <ipython-input-5-19d578e3c582>:19(<module>)
    29998    0.143    0.000    0.143    0.000 {_bisect.insort_left}
        1    0.016    0.016    0.159    0.159 <ipython-input-5-19d578e3c582>:15(<module>)

Answer 1

如果[1,3,3,3,3,3,6] 产生五个pairs (k=2)，您可以使用numpy 的broadcasting 功能 来消除Python for循环。

import numpy as np
import random

a = random.sample(xrange(1, 29999), 29998)
a = np.array(a)
# or just a = np.random.randint(1, 29999, 29998)

k = 2

创建一个新数组，其中包含所有可以组成对

的整数

b = a + k

通过在a 之间广播b 创建一个布尔数组：这会产生一个二维数组，其中包含True 的任何地方都有一个对。

c = a[:, np.newaxis] == b

总结所有True的

np.sum(c)

或者只是：

np.sum(a[:, np.newaxis] == b)

---

如果如示例 input 所示，列表仅包含唯一值，numpy 解决方案将是：

a = random.sample(xrange(1, 29999), 29998)
k = 2
a = np.array(a)
b = a + k
result = np.sum(np.in1d(b, a, assume_unique=True))

哪个更快。

事实上，如果值不是唯一的，numpy.in1d 比上面的广播解决方案快得多。通过切换参数的顺序，您计算五对 [1,3,3,3,3,3,6]。

result = np.sum(np.in1d(a, b))

---

现在有点吃乌鸦的东西：将列表变成一个集合（假设唯一值），纯 Python 解决方案比 numpy 解决方案更快。

q = 10000
a = random.sample(xrange(1, q), q-1)
a = set(a)
result = sum(n+k in a for n in a)

使用sum 消耗generator expression 不需要制作任何中间对象——这可能是其速度/效率的原因之一。

Answer 2

使用更好的算法。在你的代码中

for number in sorted_array:
    the_pair = number + k
    #check if the number+K is in the array
    if the_pair in sorted_array:
        pairs_counter += 1

您正在检查整个数组中的the_pair，因此对列表进行排序没有任何收获。由于所有元素都是整数，所以列表排序后`the_pair，如果它出现在列表中只能出现在接下来的两个位置之一。尝试类似

for index, number in sorted_array:
    if number+k in sorted_array[index+1:index+k+1]:
        <do whatever>

Answer 3

@saulspatz 是正确的，您在代码中使排序无关紧要，但是我建议您跳过排序而不是生成数千个列表切片。如果您与不可变类型（例如：tuple()）进行比较，in 操作实际上非常快。因此，我将提出以下代码：

#generate random numbers
import bisect
import random
n_integers = tuple(random.sample(xrange(1, 29999), 29998))
####cProfile
import cProfile
pr = cProfile.Profile()
pr.enable()

#the difference between numbers we are looking for
k = 2
pairs_counter = 0

#iterate over the array and calculate (number + K)
for number in n_integers:
    the_pair = number + k
    #check if the number+K is in the array
    if the_pair in n_integers:
        pairs_counter += 1

print pairs_counter

#Close cProfile
pr.disable()
pr.print_stats(sort = 'time')

输出：

29996
     1 function calls in 0.000 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    1    0.000    0.000    0.000    0.000 {method 'disable' of'_lsprof.Profiler' objects}