Created
October 3, 2013 19:24
-
-
Save landau/6815566 to your computer and use it in GitHub Desktop.
Get the k smallest ints from an array.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment