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

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

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

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 2

@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}

Answer 3

如果 [1,3,3,3,3,3,6] 产生五个对 (k=2)，您可以使用 numpy 的 broadcasting feature 消除 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 比上面的广播解决方案快得多。通过切换参数的顺序，您 count 五对 [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 不需要制作任何中间对象 - 这可能是其 speed/efficiency.

的原因之一

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

Optimizing for loop in Python to work faster

python

optimization

cprofile