Get the k smallest ints from an array.

ksmallest.py

import heapq

def heappush(heap, val):
    heapq.heappush(heap, -val)

def heappop(heap):
    return -heapq.heappop(heap)

def heapify(arr):
    for i in xrange(len(arr)):
        arr[i] = -arr[i]
    heapq.heapify(arr)

def heapsort(heap):
    arr = [None] * len(heap)
    i = len(heap) - 1
    while len(heap):
        arr[i] = heappop(heap)
        i -= 1

    return arr

# TODO custom comparator
def get_smallest(k, arr, cmp=None):
    heap = []
    heapify(heap)

    for val in arr:
        if len(heap) < k:
            heappush(heap, val)

        elif -val > heap[0]:
            heappop(heap)
            heappush(heap, val)

    return heap

if __name__ == '__main__':
    import unittest
    from random import shuffle

    class Test(unittest.TestCase):
        def setUp(self):
            self.arr = range(0, 1000)
            shuffle(self.arr)

        def test_returns_array(self):
            arr = get_smallest(5, self.arr)
            self.assertTrue(isinstance(arr, list))

        def test_5(self):
            arr = get_smallest(5, self.arr)
            self.assertEqual(len(arr), 5)

        def test_is_smallest(self):
            arr = get_smallest(5, self.arr)
            # we know 0-4 is the smallest
            filtered = filter(lambda a: a > 4, self.arr)

            for val in filtered:
                for cmp in arr:
                    self.assertTrue(cmp < val)

            self.assertEqual(heapsort(arr), range(0, 5))

    unittest.main()