python查找第k小元素代码分享
# -*- 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
def rand_select(a, k):
"""默认随机划分主元方式，k in [1, len(a)]
e[t(n)] = o(n)
"""
return _select(a, 0, len(a) - 1, k);
def _median(a, l, r):
"""对a[l..r]插入排序（原地）后选取其中位数位置"""
for j in xrange(l, r + 1):
k = a[j]
i = j
while i > l and a[i-1] > k:
a[i] = a[i-1]
i -= 1
a[i] = k
return l + int((r - l) * 0.5) # 下中位数
def _medianofmedians(a, l, r):
"""中位数的中位数方式：
1. 划分为floor(n/5)个5元组，剩下(n%5)组成最后一组。
2. 找出ceil(n/5)个组各自的中位数。先对每组插入排序，再从中选出中位数。
3. 对第2步中找出的ceil(n/5)个中位数重复上述操作，直到仅有一个中位数。
"""
if l == r:
return l
n = r - l + 1 # 元素个数
m = int(ceil(n / 5.0)) # 划分组数，每组5个元素
for i in xrange(m):
# 每组起始位和结束位
sub_l = l + i * 5
sub_r = sub_l + 4
if sub_r > r:
sub_r = r
# 对每组元素插入排序后，选取中位数
sub_m = _median(a, sub_l, sub_r) # 中位数索引
# 交换中位数到前几位
j = l + i
a[j], a[sub_m] = a[sub_m], a[j]
return _medianofmedians(a, l, l + m - 1) # 中位数的中位数
def bfprt_select(a, k):
"""中位数的中位数方式（bfprt算法）
t(n) = o(n)
"""
return _select(a, 0, len(a) - 1, k, _medianofmedians);
def _median3(a, l, r):
"""三数中位数方式，取l,r,(l+r)/2三数中位数"""
c = (l + r) / 2
keys = [l, c, r]
i = _median(keys, 0, 2)
return keys[i]
def median_select(a, k):
"""三数中位数方式，以消除最坏情况"""
return _select(a, 0, len(a) - 1, k, _median3);
if __name__ == '__main__':
import random, time
from copy import copy
print('preparing data...')
n = 1000000
nums = range(n)
random.shuffle(nums)
print('ready go!')
def timeit(fnc, *args, **kargs):
print('%s starts processing' % fnc.__name__)
begtime = time.clock()
retval = fnc(*args, **kargs)
endtime = time.clock()
print('%s takes time : %f' % (fnc.__name__, endtime - begtime))
return retval
test_methods = [rand_select, bfprt_select, median_select]
k = random.randrange(n) + 1
dashes = '---' * 10
for test in test_methods:
print(dashes)
nums_new = copy(nums)
result = timeit(test, nums_new, k)
print('the %dth smallest element: %d' % (k, result))