radix sort in python

radix_sort.py

#coding:UTF-8

from random import randint
from functools import partial
from counting_sort import counting_sort

'''
for simplicity: use base 10
b - base
k - range of values of keys of the input
O((n+b)*lgb k)
when b = n, the above expression is minimized => O(n*lgn k)
'''

def radix_sort(arr, n, k):
    def get_int_pos(i, p):
        try:
            return int(str(i)[-p])
        except IndexError:
            return 0

    sorted_arr = arr
    for i in range(1, len(str(k))+1): # sort on each digit
        # *this step must be stable*
        sorted_arr = counting_sort(sorted_arr, n, k, key=partial(get_int_pos, p=i))
    return sorted_arr

if __name__ == '__main__':
    n = 10
    k = 1000
    arr = [randint(0, k) for i in range(n)]
    sorted_arr = radix_sort(arr, n, k)
    print("arr: {}".format(arr))
    print("sorted_arr: {}".format(sorted_arr))