python查找第k小元素代码分享

771次阅读  |  发布于5年以前

复制代码 代码如下:

-- 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))  

Copyright© 2013-2020

All Rights Reserved 京ICP备2023019179号-8