代码如下:
# -*- coding: utf-8 -*-
from random import randint
from math import ceil, floor
def _partition(A, l, r, i):
"""以A[i]为主元划分数组A[l..r],使得:
A[l..m-1] <= A[m] < A[m+1..r]
"""
A[i], A[r] = A[r], A[i] # i交换到末位r,作为主元
pivot = A[r] # 主元
m = l # 索引标记
for n in xrange(l, r): # l..r-1
if A[n] <= pivot:
A[m], A[n] = A[n], A[m] # 交换
m += 1 # 后移
A[m], A[r] = A[r], A[m] # 主元到m位
return m
def _rand(A, l, r):
"""随机划分主元"""
return randint(l, r) # A[l..r]随机取一个
def _select(A, l, r, k, pivot_selector = _rand):
"""利用快排,得A[l..r]中第k小的数,k in [l+1,r+1]:
其尾递归方式,伪码如下:
SELECT(A, l, r, k)
1 while true:
2 i ← ? // 划分主元位置
3 m ← PARTITION(A, l, r, i) // 数组划分
4 n ← m - l + 1 // A[l..m]元素个数
5 if k = n // 检查A[m]是否是第k小的元素
6 then return A[m]
7 elseif k < n // 左划分区
8 r = m - 1
9 else // 右划分区
10 k = k - n
11 l = m + 1
Args:
pivot_selector(Function): 主元选取方法,默认随机方式
"""
if not A:
return None
if l == r:
return A[l]
while True:
i = pivot_selector(A, l, r)
m = _partition(A, l, r, i)
n = m - l + 1
if k == n:
return A[m]
elif k < n:
r = m - 1
else:
k = k - n
l = m + 1