another solution in problems.py