fork download
  1. import random
  2. import time
  3.  
  4. def median5(a): #find median size<=5
  5. n=len(a)
  6. return sorted(a)[n//2]
  7.  
  8.  
  9. def kth(a,k):
  10. if k==1:
  11. return min(a)
  12. n=len(a)
  13.  
  14. median=[]
  15. i=0
  16. while i<n//5:
  17. median.append(median5(a[i*5:i*5+5]))
  18. i+=1
  19. if n%5:
  20. median.append(median5(a[i*5:]))
  21.  
  22. if len(median)==1:
  23. pivot=a[0]
  24. else:
  25. pivot=kth(median,len(median)//2)
  26.  
  27. less=filter(lambda x:x<pivot,a)
  28. nEq=a.count(pivot)
  29. greater=filter(lambda x:x>pivot,a)
  30. nL=len(less)
  31. nG=len(greater)
  32. if k<=nL:
  33. return kth(less,k)
  34. elif k<=nL+nEq:
  35. return pivot
  36. return kth(greater,k-nL-nEq)
  37.  
  38.  
  39.  
  40. n=1000000
  41. a=[]
  42. for i in range(n):
  43. a.append(random.randint(0,n))
  44.  
  45. k=random.randint(1,n)
  46. print kth(a,k),sorted(a[:])[k-1]
  47.  
Success #stdin #stdout 2.89s 11560KB
stdin
Standard input is empty
stdout
490580 490580