我正在尝试在 Python 中使用 numba 实现快速排序算法。
它似乎比 numpy 排序函数慢很多。
我该如何改进它？我的代码在这里:
import numba as nb
@nb.autojit
def quick_sort(list_):
"""
Iterative version of quick sort
"""
#temp_stack = []
#temp_stack.append((left,right))
max_depth = 1000
left = 0
right = list_.shape[0]-1
i_stack_pos = 0
a_temp_stack = np.ndarray( ( max_depth, 2), dtype=np.int32 )
a_temp_stack[i_stack_pos,0] = left
a_temp_stack[i_stack_pos,1] = right
i_stack_pos+=1
#Main loop to pop and push items until stack is empty
while i_stack_pos>0:
i_stack_pos-=1
right = a_temp_stack[ i_stack_pos, 1 ]
left = a_temp_stack[ i_stack_pos, 0 ]
piv = partition(list_,left,right)
#If items in the left of the pivot push them to the stack
if piv-1 > left:
#temp_stack.append((left,piv-1))
a_temp_stack[ i_stack_pos, 0 ] = left
a_temp_stack[ i_stack_pos, 1 ] = piv-1
i_stack_pos+=1
#If items in the right of the pivot push them to the stack
if piv+1 < right:
a_temp_stack[ i_stack_pos, 0 ] = piv+1
a_temp_stack[ i_stack_pos, 1 ] = right
i_stack_pos+=1
@nb.autojit( nopython=True )
def partition(list_, left, right):
"""
Partition method
"""
#Pivot first element in the array
piv = list_[left]
i = left + 1
j = right
while 1:
while i <= j and list_[i] <= piv:
i +=1
while j >= i and list_[j] >= piv:
j -=1
if j <= i:
break
#Exchange items
list_[i], list_[j] = list_[j], list_[i]
#Exchange pivot to the right position
list_[left], list_[j] = list_[j], list_[left]
return j
我的测试代码在这里:
x = np.random.random_integers(0,1000,1000000)
y = x.copy()
quick_sort( y )
z = np.sort(x)
np.testing.assert_array_equal( z, y )
y = x.copy()
with Timer( 'nb' ):
numba_fns.quick_sort( y )
with Timer( 'np' ):
x = np.sort(x)
更新:
我重写了函数以强制代码的循环部分在 nopython 模式下运行。 while 循环似乎没有导致 nopython 失败。但是，我没有获得任何性能提升:
@nb.autojit
def quick_sort2(list_):
"""
Iterative version of quick sort
"""
max_depth = 1000
left = 0
right = list_.shape[0]-1
i_stack_pos = 0
a_temp_stack = np.ndarray( ( max_depth, 2), dtype=np.int32 )
a_temp_stack[i_stack_pos,0] = left
a_temp_stack[i_stack_pos,1] = right
i_stack_pos+=1
#Main loop to pop and push items until stack is empty
return _quick_sort2( list_, a_temp_stack, left, right )
@nb.autojit( nopython=True )
def _quick_sort2( list_, a_temp_stack, left, right ):
i_stack_pos = 1
while i_stack_pos>0:
i_stack_pos-=1
right = a_temp_stack[ i_stack_pos, 1 ]
left = a_temp_stack[ i_stack_pos, 0 ]
piv = partition(list_,left,right)
#If items in the left of the pivot push them to the stack
if piv-1 > left:
a_temp_stack[ i_stack_pos, 0 ] = left
a_temp_stack[ i_stack_pos, 1 ] = piv-1
i_stack_pos+=1
if piv+1 < right:
a_temp_stack[ i_stack_pos, 0 ] = piv+1
a_temp_stack[ i_stack_pos, 1 ] = right
i_stack_pos+=1
@nb.autojit( nopython=True )
def partition(list_, left, right):
"""
Partition method
"""
#Pivot first element in the array
piv = list_[left]
i = left + 1
j = right
while 1:
while i <= j and list_[i] <= piv:
i +=1
while j >= i and list_[j] >= piv:
j -=1
if j <= i:
break
#Exchange items
list_[i], list_[j] = list_[j], list_[i]
#Exchange pivot to the right position
list_[left], list_[j] = list_[j], list_[left]
return j
请您参考如下方法:
一个可能有帮助的小建议(但正如您在问题的评论中正确地被告知的那样，您将难以击败纯 C 实现):
您要确保其中大部分是在“nopython”模式下完成的 (@jit(nopython=True))。在你的函数之前添加它，看看它在哪里中断。同时在您的函数上调用 inspect_types()，看看它是否能够正确识别它们。
在您的代码中，可能会迫使它进入对象模式(与 nopython 模式相反)的一件事是分配一个 numpy 数组。尽管 numba 可以在 nopython 模式下单独编译循环，但我不知道它是否可以为 while 循环做到这一点。调用 inspect_types 会告诉您。
创建 numpy 数组同时确保其余数组处于 nopython 模式的通常工作方法是创建一个包装函数。
@nb.jit(nopython=True) # make sure it can be done in nopython mode
def _quick_sort_impl(list_,output_array):
...most of your code goes here...
@nb.jit
def quick_sort(list_):
# this code won't compile in nopython mode, but it's
# short and isolated
max_depth = 1000
a_temp_stack = np.ndarray( ( max_depth, 2), dtype=np.int32 )
_quick_sort_impl(list_,a_temp_stack)
