Extracted bitmap_sort function

column1/problem3.py

@@ -1,15 +1,16 @@
 from bitarray import bitarray
 
+def bitmap_sort(nums, n):
+    arr = (n + 1) * bitarray('0')
+    for num in nums:
+        arr[num] = 1
+
+    return (i for (i, bit) in enumerate(arr) if bit)
+
 if __name__ == '__main__':
     import sys
     n = int(sys.argv[1])
-    arr = (n + 1) * bitarray('0')
+    nums = (int(line) for line in sys.stdin)
 
-    for line in sys.stdin:
-        num = int(line)
-        arr[num] = 1
-
-    for (i, bit) in enumerate(arr):
-        if bit:
-            print i
-        
+    for i in bitmap_sort(nums, n):
+        print i