查找第K小的元素的O(N)算法

话说这个问题，比较挫的解决方案有

1.先排序，然后找到第K小的，复杂度是O(nlogn)

2.选择排序来搞，选择排序是O(kn)，

3.堆排序是O(nlogk)

4.比较好的解决方案是利用类似快速排序的划分思想来找到第K小，复杂度为O(n)，但是最坏情况可能达到O(n^2)

5.还有种方法可以使得最坏情况也是O(n)。

我们先来看用快速排序的思想来搞的方案。快速排序是找到一个数，然后把所有数分为小于等于那个数的一堆，和大于那个数的一堆，然后两段分别递归来排序，而我们查找算法里，由于知道第K小的元素会在哪一堆，这样只需要递归其中一对即可。

 

1. import random

2.  

3. def partition(arr, left, right, pivot):

4.     v = arr[pivot]

5.     arr[pivot], arr[right-1] = arr[right-1], arr[pivot]

6.     index = left

7.     for i in xrange(left, right):

8.         if arr[i] <= v:

9.             arr[i], arr[index] = arr[index], arr[i]

10.             index += 1

11.     return index-1

12.  

13. def select(arr, left, right, k):

14.     while right - left > 1:

15.         index = partition(arr, left, right, random.randint(left, right-1))

16.         dist = index - left + 1

17.         if dist == k:

18.             return arr[index]

19.         if dist < k:

20.             k -= dist

21.             left = index + 1

22.         else:

23.             right = index

24.     return arr[left]

 

之后arr是要查找的数组，调用select即可找到第K小元素，如果pivot元素选的不好那么这个算法最坏的情况是O(n^2)。

现在讨论最坏情况下也是O(n)的方案，把所有的数分为5个一堆，那么总共会有n/5堆，对于每堆我们可以很快的找到中位数（因为只有5个所以很容易嘛），之后调用当前算法找到这n/5个中位数的中位数，用这个数来做pivot，所以这个算法被叫做Median of Medians algorithm。

把中位数的中位数作为pivot的话，那么原数组中便会有3/5*1/2个也就是3/10个小于等于这个pivot的，同理会有3/10大于这个pivot的，所以最坏情况下，数组被分为30%，70%或者70%，30%的两部分。

T(n)<=T(n/5)+T(7/10*n)+O(n)<=c*n*(1+9/10+(9/10)^2....) 
所以T(n)=O(n)

也就是最坏情况下是O(n)。

 

1. import heapq

2.  

3. def partition(arr, left, right, pivot):

4.     v = arr[pivot]

5.     arr[pivot], arr[right-1] = arr[right-1], arr[pivot]

6.     index = left

7.     for i in xrange(left, right):

8.         if arr[i] <= v:

9.             arr[i], arr[index] = arr[index], arr[i]

10.             index += 1

11.     return index-1

12.  

13. def select_heap(arr, left, right, k):

14.     tmp = arr[left:right]

15.     heapq.heapify(tmp)

16.     [heapq.heappop(tmp) for i in xrange(k-1)]

17.     return heapq.heappop(tmp)

18.  

19. def median(arr, left, right):

20.     num = (right - left - 1) / 5

21.     for i in xrange(num+1):

22.         sub_left = left + i*5

23.         sub_right = sub_left + 5

24.         if sub_right > right:

25.             sub_right = right

26.         m_index = select_heap(arr, sub_left, sub_right, (sub_right-sub_left)/2)

27.         arr[left+i], arr[m_index] = arr[m_index], arr[left+i]

28.     return select(arr, left, left+num+1, (num+1)/2)

29.  

30. def select(arr, left, right, k):

31.     while right - left > 1:

32.         pivot = median(arr, left, right)

33.         index = partition(arr, left, right, pivot)

34.         dist = index - left + 1

35.         if dist == k:

36.             return arr[index]

37.         if dist < k:

38.             k -= dist

39.             left = index + 1

40.         else:

41.             right = index

42.     return arr[left]

 

同理，如果快速排序每次选pivot时用Median of Medians algorithm也可以把最坏情况降低为O(nlogn)的。