import math def solution(w, h, s): return str(sum([s**cycle for cycle in generate_cycles(w, h)])/(math.factorial(w)*math.factorial(h))) def generate_cycles(w, h): yield 1 tests = [ ([2, 3, 4], '430'), ([2, 2, 2], '7'), ([1, 1, 2], '2'), ([1, 1, 3], '3'), ([2, 1, 2], '3'), ([1, 2, 2], '3'), ([3, 1, 2], '4'), ([1, 3, 2], '4'), ([2, 3, 2], '13'), ] for i, o in tests: result = solution(*i) print (i, result == o, result, o)