problem source https www hackerrank comchallengesthepowersumproblem find the number of ways that a given integer x can be expressed as the sum of the nth powers of unique natural numbers for example if x13 and n2 we have to find all combinations of unique squares adding up to 13 the only solution is 2232 constraints 1x1000 2n10 backtrack13 2 1 0 0 0 1 backtrack10 2 1 0 0 0 1 backtrack10 3 1 0 0 0 0 backtrack20 2 1 0 0 0 1 backtrack15 10 1 0 0 0 0 backtrack16 2 1 0 0 0 1 backtrack20 1 1 0 0 0 64 if the sum of the powers is equal to neededsum then we have a solution if the sum of the powers is less than neededsum then continue adding powers if the power of i is less than neededsum then try with the next power solve13 2 1 solve10 2 1 solve10 3 0 solve20 2 1 solve15 10 0 solve16 2 1 solve20 1 traceback most recent call last valueerror invalid input neededsum must be between 1 and 1000 power between 2 and 10 solve10 5 traceback most recent call last valueerror invalid input neededsum must be between 1 and 1000 power between 2 and 10 backtrack 13 2 1 0 0 0 1 backtrack 10 2 1 0 0 0 1 backtrack 10 3 1 0 0 0 0 backtrack 20 2 1 0 0 0 1 backtrack 15 10 1 0 0 0 0 backtrack 16 2 1 0 0 0 1 backtrack 20 1 1 0 0 0 64 if the sum of the powers is equal to needed_sum then we have a solution if the sum of the powers is less than needed_sum then continue adding powers if the power of i is less than needed_sum then try with the next power solve 13 2 1 solve 10 2 1 solve 10 3 0 solve 20 2 1 solve 15 10 0 solve 16 2 1 solve 20 1 traceback most recent call last valueerror invalid input needed_sum must be between 1 and 1000 power between 2 and 10 solve 10 5 traceback most recent call last valueerror invalid input needed_sum must be between 1 and 1000 power between 2 and 10 return the solutions_count
def backtrack(
    needed_sum: int,
    power: int,
    current_number: int,
    current_sum: int,
    solutions_count: int,
) -> tuple[int, int]:
    
    if current_sum == needed_sum:
        
        solutions_count += 1
        return current_sum, solutions_count

    i_to_n = current_number**power
    if current_sum + i_to_n <= needed_sum:
        
        current_sum += i_to_n
        current_sum, solutions_count = backtrack(
            needed_sum, power, current_number + 1, current_sum, solutions_count
        )
        current_sum -= i_to_n
    if i_to_n < needed_sum:
        
        current_sum, solutions_count = backtrack(
            needed_sum, power, current_number + 1, current_sum, solutions_count
        )
    return current_sum, solutions_count


def solve(needed_sum: int, power: int) -> int:
    
    if not (1 <= needed_sum <= 1000 and 2 <= power <= 10):
        raise ValueError(
            "Invalid input\n"
            "needed_sum must be between 1 and 1000, power between 2 and 10."
        )

    return backtrack(needed_sum, power, 1, 0, 0)[1]  


if __name__ == "__main__":
    import doctest

    doctest.testmod()