power_sum.py

"""
Problem source: https://www.hackerrank.com/challenges/the-power-sum/problem
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 X=13 and N=2.
We have to find all combinations of unique squares adding up to 13.
The only solution is 2^2+3^2. Constraints: 1<=X<=1000, 2<=N<=10.
"""


defbacktrack(
needed_sum: int,
power: int,
current_number: int,
current_sum: int,
solutions_count: int,
) ->tuple[int, int]:
"""
 >>> 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)
 """
ifcurrent_sum==needed_sum:
# If the sum of the powers is equal to needed_sum, then we have a solution.
solutions_count+=1
returncurrent_sum, solutions_count

i_to_n=current_number**power
ifcurrent_sum+i_to_n<=needed_sum:
# If the sum of the powers is less than needed_sum, then continue adding powers.
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
ifi_to_n<needed_sum:
# If the power of i is less than needed_sum, then try with the next power.
current_sum, solutions_count=backtrack(
needed_sum, power, current_number+1, current_sum, solutions_count
 )
returncurrent_sum, solutions_count


defsolve(needed_sum: int, power: int) ->int:
"""
 >>> 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.
 """
ifnot (1<=needed_sum<=1000and2<=power<=10):
raiseValueError(
"Invalid input\n"
"needed_sum must be between 1 and 1000, power between 2 and 10."
 )

returnbacktrack(needed_sum, power, 1, 0, 0)[1] # Return the solutions_count


if__name__=="__main__":
importdoctest

doctest.testmod()