Задача о количестве цифр в записи чисел от 1 до n

digits.py

def get_digit_count(n):
	digit_count = [0] * 10
	s = n
	div = 1
	while s > 0:
		s //= 10
		x = n - (div - 1)
		full_period = x // (10 * div)
		unfull = x % (10 * div)
		unfull_periods = unfull // div
		unfull_left = unfull % div
		digit_count = [i + full_period * div for i in digit_count]

		t = 0
		if unfull_periods > 0:
			for t in range(1, unfull_periods+1):
				digit_count[t] += div
		digit_count[(t + 1) % 10] += unfull_left
	
		div *= 10
		
	
	return digit_count
	
def get_digit_count_bad(n):
	digit_count = [0] * 10
	for i in range(1, n + 1):
		s = i
		while s > 0:
			digit_count[s % 10] += 1
			s //= 10
	return digit_count
	

def main():
	print(get_digit_count(1058))
	print(get_digit_count_bad(1058))

main()