The Python implementation of getting n largest elements in unsorted array.

nlargest.py

'''
The implementation of getting n largest elements in the array.
The algorithms follows the page:
https://www.geeksforgeeks.org/k-largestor-smallest-elements-in-an-array/
'''
import operator
from math import inf


def nlargest_bubble(n, items):
    items = items[:]
    for _ in range(min(n, len(items))):
        for i in range(len(items) - 1):
            if items[~i] > items[~i-1]:
                items[~i], items[~i-1] = items[~i-1], items[~i]
    return items[:n]


def nlargest_narray(n, items):
    if n >= len(items):
        return sorted(items, reverse=True)
    items = items[:]
    for i in range(n, len(items)):
        nth, key = min((items[key], key) for key in range(n))
        if items[i] > nth:
            items[key], items[i] = items[i], items[key]
    return sorted(items[:n], reverse=True)


def nlargest_sort(n, items):
    return sorted(items, reverse=True)[:n]


def heapify(arr, cur, beg=0, end=None, opr=operator.gt):
    end = len(arr) if end is None else end
    mid = (beg + end) >> 1
    off = beg - 1
    while cur < mid:
        tmp = cur - off
        l, r = (tmp << 1) + off, (tmp << 1) + off + 1
        if r < end and opr(arr[r], arr[l]) and opr(arr[r], arr[cur]):
            arr[cur], arr[r] = arr[r], arr[cur]
            cur = r
        elif opr(arr[l], arr[cur]):
            arr[cur], arr[l] = arr[l], arr[cur]
            cur = l
        else:
            break


def max_heapify(arr, cur, beg=0, end=None):
    return heapify(arr, cur, beg, end, opr=operator.gt)


def min_heapify(arr, cur, beg=0, end=None):
    return heapify(arr, cur, beg, end, opr=operator.lt)


def build_heap(arr, beg=0, end=None, opr=operator.gt):
    end = len(arr) if end is None else end
    mid = (beg + end) >> 1
    for i in range(mid - 1, -1, -1):
        heapify(arr, i, beg, end, opr)


def build_max_heap(arr, beg=0, end=None):
    end = len(arr) if end is None else end
    return build_heap(arr, beg, end, opr=operator.gt)


def build_min_heap(arr, beg=0, end=None):
    end = len(arr) if end is None else end
    return build_heap(arr, beg, end, opr=operator.lt)


def nlargest_max_heap(n, items):
    items = items[:]
    build_max_heap(items)
    res = []
    for _ in range(min(n, len(items))):
        res.append(items[0])
        items[0] = -inf
        max_heapify(items, 0)
    return res


def nlargest_oder_statistics(n, items):
    items = items[:]
    k = len(items)
    beg, end = 0, k
    while end - beg > 1:
        # find the median
        medians = items[beg:end]
        while len(medians) > 1:
            parts = [sorted(medians[i:i+5]) for i in range(0, len(medians), 5)]
            medians = [part[len(part) >> 1] for part in parts]
        median = medians[0]
        # partition
        l = beg
        for r in range(beg, end-1):
            if items[r] == median:
                items[end-1], items[r] = items[r], items[end-1]
            if items[r] < median:
                items[l], items[r] = items[r], items[l]
                l += 1
        items[end-1], items[l] = items[l], items[end-1]
        if k - n < l:
            end = l
        elif k - n > l:
            beg = l + 1
        else:
            break
    return sorted(items[-n:], reverse=True)


def nlargest_min_heap(n, items):
    if n < len(items):
        items = items[:]
        build_min_heap(items, 0, n)
        for i in range(n, len(items)):
            if items[i] > items[0]:
                items[0], items[i] = items[i], items[0]
                min_heapify(items, 0, 0, n)
    return sorted(items[:n], reverse=True)


if __name__ == '__main__':
    from random import shuffle
    nums = list(range(1000))
    shuffle(nums)

    for n in (15, 500, 1000, 1200):
        ans = sorted(nums, reverse=True)[:n]
        assert nlargest_bubble(n, nums) == ans
        assert nlargest_narray(n, nums) == ans
        assert nlargest_sort(n, nums) == ans
        assert nlargest_max_heap(n, nums) == ans
        assert nlargest_oder_statistics(n, nums) == ans
        assert nlargest_min_heap(n, nums) == ans