Skip to content

Instantly share code, notes, and snippets.

@stepchowfun
Last active February 1, 2017 21:38
Show Gist options
  • Save stepchowfun/7233877 to your computer and use it in GitHub Desktop.
Save stepchowfun/7233877 to your computer and use it in GitHub Desktop.
Worst-case linear-time selection algorithm in Python. Note that in practice there are faster ways to find the k-th largest element in a list, even if this implementation is asymptotically faster.
#!/usr/bin/python -O
# partition A[p:r] in place about x, and return the final position of the pivot
def partition(A, p, r, x):
# find the index of the pivot
i = -1
for j in range(p, r):
if A[j] == x:
i = j
break
# move the pivot to the end
if i != -1:
t = A[r - 1]
A[r - 1] = A[i]
A[i] = t
# keep track of the end of the "less than" sublist
store_index = p
# iterate
for j in range(p, r - 1):
if A[j] < x:
t = A[j]
A[j] = A[store_index]
A[store_index] = t
store_index += 1
# put the pivot in its final place
if i != -1:
t = A[r - 1]
A[r - 1] = A[store_index]
A[store_index] = t
# return the store index
return store_index
# find the ith biggest element in A[p:r]
def select(A, p, r, i):
# make a copy of the array
A = A[:]
# divide the n elements of A into n / 5 groups
groups = [[]] * (((r - p) + 4) / 5)
for x in range(p, r):
groups[(x - p) / 5] = groups[(x - p) / 5] + [A[x]]
# find the median of each group
medians = [sorted(l)[(len(l) - 1) / 2] for l in groups]
# find the median of medians
if len(medians) == 1:
median_to_rule_them_all = medians[0]
else:
median_to_rule_them_all = select(medians, 0, len(medians), (len(medians) - 1) / 2)
# partition A around the median of medians
partition_index = partition(A, p, r, median_to_rule_them_all)
# base case
if p + i < partition_index:
return select(A, p, partition_index, i)
if p + i > partition_index:
return select(A, partition_index + 1, r, p + i - partition_index - 1)
return A[p + i]
@BazzalSeed
Copy link

the other algorithm you were mentioning is quick select


def quickselect(alist, k):
    start, end = 0, len(alist) - 1
    return quickselecthelper(alist, start, end, k)


def quickselecthelper(alist, start, end, k):
    if start <= end:
        split = random_partition(alist, start, end)
        print start + k, split, start, end
        if k == split:
            print "found"
            return alist[start + k]
        elif k < split:
            return quickselecthelper(alist, start, split - 1, k)
        else:
            return quickselecthelper(alist, split + 1, end, k)


def random_partition(alist, start, end):
    from random import randint
    pivot = randint(start, end)
    temp = alist[start]
    alist[start] = alist[pivot]
    alist[pivot] = temp

    leftmark = start + 1
    rightmark = end
    done = False
    pivotvalue = alist[start]

    while not done:
        while leftmark <= rightmark and alist[leftmark] <= pivotvalue:
            leftmark += 1
        while rightmark >= leftmark and alist[rightmark] >= pivotvalue:
            rightmark -= 1
        if leftmark > rightmark:
            done = True
        else:
            swap(alist, leftmark, rightmark)

    swap(alist, start, rightmark)
    return rightmark


def swap(alist, a, b):
    temp = alist[a]
    alist[a] = alist[b]
    alist[b] = temp


def random_partition(alist, start, end):
    from random import randint
    pivot = randint(start, end)
    swap(alist, start, pivot)

    leftmark = start + 1
    rightmark = end
    done = False
    pivotvalue = alist[start]

    while not done:
        while leftmark <= rightmark and alist[leftmark] <= pivotvalue:
            leftmark += 1
        while rightmark >= leftmark and alist[rightmark] >= pivotvalue:
            rightmark -= 1
        if leftmark > rightmark:
            done = True
        else:
            swap(alist, leftmark, rightmark)

    swap(alist, start, rightmark)
    return rightmark

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment