counting sort in Python

counting_sort.py

#coding:UTF-8

from random import randint
from math import inf

'''
running time: O(n+k)
'''

def counting_sort(arr, n, k, key=lambda x: x):
    counting_arr = [[] for i in range(k+1)] # O(k)
    sorted_arr = [] # O(1)
    for i in arr: # O(n)
        counting_arr[key(i)].append(i) # O(1) => stable sort
    for c in counting_arr: # O(k)
        sorted_arr.extend(c) # O(len(c))
    return sorted_arr


def counting_sort_v2(arr, n, k, key=lambda x: x):
    counting_arr = [0 for i in range(k+1)] # O(k)
    sorted_arr = [inf for i in range(len(arr))] # O(1)

    for i in arr: # O(n)
        counting_arr[key(i)] += 1 # O(1)

    # counting_arr => position_arr
    sum = 0
    for i, v in enumerate(counting_arr):
        counting_arr[i] = sum
        sum += v

    for i in arr:
        idx = counting_arr[key(i)]
        sorted_arr[idx] = i
        counting_arr[key(i)] += 1

    return sorted_arr


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