查找第K小的元素

from http://www.isnowfy.com/top-k-number/

感觉这是个经典问题了，但是今天看维基百科的时候还是有了新的发现，话说这个问题，比较挫的解决方案有先排序，然后找到第K小的，复杂度是O(nlogn)，还有就是利用选择排序或者是堆排序来搞，选择排序是O(kn)，堆排序是O(nlogk)，比较好的解决方案是利用类似快速排序的思想来找到第K小，复杂度为O(n)，但是最坏情况可能达到O(n^2)，不过今天要说的，就是还有种方法可以使得最坏情况也是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)的。