Partition function with itertools
The elephants in the room:


Single letter variable names are to be avoided. Each variable should have a suggestive meaningful name.
They probably expected you to write automated tests, just printing a result for visual inspection is not a real way of testing. assert is easy to use and can get you started on automated testing.