Be careful. You need to leave the function signature just as it is provided, because otherwise our assessment engine will not be able to run your program. With that in mind, feel free to add helper functions, imports/includes...
(In some tasks there are extra restrictions on how much code you are allowed to change - the goal is to understand the code and make a minimal edit that fixes the given problem.)
Can I modify the initial code?
