secret
Created

  • Download Gist
nsmallest.py
Python
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
def partition(A, idx, start=0, end=None):
if end is None:
end = len(A)
pivot = A[idx]
A[idx], A[end - 1] = A[end - 1], A[idx]
pos = start
for i in range(start, end):
if A[i] < pivot:
if i != pos:
A[pos], A[i] = A[i], A[pos]
pos += 1
A[pos], A[end - 1] = A[end - 1], A[pos]
return pos
 
def nsmallest(A, n, start=0, end=None):
if end is None:
end = len(A)
partition_idx = partition(A, start, start, end)
if partition_idx == n:
return A[:n]
elif partition_idx < n:
return nsmallest(A, n, partition_idx + 1, end)
else:
return nsmallest(A, n, start, partition_idx)

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.