Counting Sort Estable: Ordenar una lista de enteros con colores asignados

CountingSort.py

from termcolor import colored #pip install termcolor
import random, re

COLORS = ["green", "blue", "red", "yellow", "magenta", "cyan"]
ANSI_ESCAPE = re.compile(r'\x1b[^m]*m')

"""Utilizacion:
	Desde la terminal de Python: from CountingSort import counting_sort
	counting_sort([1,2,3])
	Procede a colorear la lista de enteros y luego ordenarla. El color sirve para mostrar la estabilidad del ordenamiento.
"""

class ColoredList(list):
	# Los colores por terminal son nada mas que cadenas llamadas "ANSI escape code".
	# Por ejemplo, si yo agrego '\x1b[32m' antes de una cadena, todo el resto saldra en verde
	# Con esa idea, esta subclase de la lista de Python recibe una lista de elementos enteros y les pega ANSI codes a la izquierda y a la derecha para poder agregarle un color a cada elemento 
	# De una lista de enteros pasa a una lista de cadenas
	def __init__(self, data=[]):
		list.__init__(self,data)
		self.color()

	def __str__(self):
		# Para imprimir correctamente hay que imprimir uno por uno. Por eso el hack del join y hacerlo parecer un print de una lista normal
		# Si bien se imprimen [1,2,3], tener en cuenta que cada elemento es una cadena.
		return '[' + ', '.join(self) + ']'

	def color(self):
		# Recibe una lista de elementos de cualquier tipo y les agrega color.
		# Modifica in place a la lista y devuelve una lista de cadenas de los mismos elementos con colores elegidos al azar.
		no_color = self.remove_color()
		for i in range(len(self)):
			self[i] = colored(no_color[i], random.choice(COLORS))
		return self

	def remove_color(self):
		# Devuelve la lista de enteros original
		# No in place
		no_color = []
		for x in range(len(self)):
			no_color.append(int(ANSI_ESCAPE.sub('', str(self[x]))))
		return no_color

def counting_sort(array):
	array = ColoredList(array)
	print('Array Original = ', array)
	no_color = array.remove_color()
	n = len(no_color)
	minimo,maximo = min(no_color), max(no_color)
	k = maximo-minimo
	contador = [0] * (k+1)
	for i in range(n):
		numero = no_color[i]
		contador[numero-minimo] += 1
	print('Contador = ', contador)
	sumatoria = [0] 
	for i in range(1,k+1):
		cantidad = contador[i-1]
		suma_actual = sumatoria[i-1]
		sumatoria.append(suma_actual + cantidad)
	print('Sumatoria Acumulada = ', sumatoria)
	resultado = ColoredList([0] * n)
	for i in range(n):
		numero = no_color[i]
		posicion = sumatoria[numero-minimo]
		sumatoria[numero-minimo] += 1
		resultado[posicion] = array[i]
	print('Resultado = ', resultado)
	print()
	
if __name__ == '__main__':
	"""De ser corrido en crudo, se ejecuta X veces, para demostrar su funcionamiento"""
	PRUEBAS_RANDOM = 2
	CANT, RANGO = 10, 10
	for x in range(PRUEBAS_RANDOM):
		array = [random.choice(range(1,RANGO)) for i in range(CANT)]
		counting_sort(array)