Find K largest elements in a given collection of numbers.

k_largest_elements.py

import sys
import random

# O(n * lg(k))
def k_largest_elements(arr, k):
  k_elems        = arr[0:k]
  k_elems_heaped = min_heapify(k_elems)
  for idx in range(k + 1, len(arr)):
    if arr[idx] > k_elems_heaped[0]:
      k_elems_heaped[0] = arr[idx]
      min_heapify_element(k_elems_heaped, 0)
  return k_elems_heaped

# O(n * lg(n))
# But the more precise calculation can prove it's actually O(n)
def min_heapify(arr):
  idx = len(arr) / 2
  while idx >= 0:
    min_heapify_element(arr, idx)
    idx -= 1
  return arr

# O(lg(n))
def min_heapify_element(arr, idx):
  left  = arr[idx * 2 + 1] if (idx * 2 + 1) < len(arr) else sys.maxint
  right = arr[idx * 2 + 2] if (idx * 2 + 2) < len(arr) else sys.maxint
  if arr[idx] <= left and arr[idx] <= right:
    return
  min_idx = idx     if arr[idx]     <= left  else idx * 2 + 1
  min_idx = min_idx if arr[min_idx] <= right else idx * 2 + 2
  swap(arr, idx, min_idx)
  return min_heapify_element(arr, min_idx)

def swap(arr, idx1, idx2):
  tmp       = arr[idx2]
  arr[idx2] = arr[idx1]
  arr[idx1] = tmp

if __name__ == "__main__":
  nums = list()
  for i in range(1000):
    nums.append(random.randint(1, 1000))

  k_largest = k_largest_elements(nums, 100)
  print(k_largest)