Reputation: 561
I am trying to write a program that takes as input:
arr
)sum
)It should return all the possible combinations of numbers in arr
that add up to sum
.
This is what I have so far:
def printAllSubsetsRec(arr, v, sum):
if sum == 0:
return [v]
if len(arr) == 0:
return
v1 = [] + v
v1.append(arr[0])
without_first = printAllSubsetsRec(arr[1:], v, sum)
with_first = printAllSubsetsRec(arr[1:], v1, sum - arr[0])
if with_first and without_first:
return with_first + without_first
elif with_first:
return with_first
elif without_first:
return without_first
def array_sums(arr, sum):
v = []
return printAllSubsetsRec(arr, v, sum)
The problem is that it doesn't return all the subsets including repetitions.
For example:
print(array_sums([1,3,5],5))
# [[1, 1, 1, 1, 1], [1, 1, 3], [1, 3, 1], [3, 1, 1], [5]]
How can I do that?
Upvotes: 5
Views: 1013
Reputation: 7867
At each recursion, I created a new branch for every number in arr
. I kept the branches whose sum matches the target, and stopped exploring branches whose sum overshoots the target.
Faster (pass the accumulated sum of a branch down the call chain)
def array_sums(arr: Set[int], target: int) -> List[List[int]]:
smallest = min(arr)
def go(
in_progress: List[Tuple[List[int], int]],
complete: List[List[int]]
) -> List[List[int]]:
now_in_progress, newly_complete = [], []
for branch, sum_ in in_progress:
for i in arr:
# anything short of `target` by less than `smallest` will overshoot
if sum_ + i <= target - smallest:
now_in_progress.append((branch + [i], sum_ + i))
elif sum_ + i == target:
newly_complete.append(branch + [i])
newly_complete += complete
return newly_complete if not now_in_progress else go(
now_in_progress, newly_complete
)
return go([([], 0)], [])
Simpler and purely functional (calculate the sum of a branch at each recursion)
def array_sums(arr: Set[int], target: int) -> List[List[int]]:
def go(acc: List[List[int]]) -> List[List[int]]:
in_progress = [
branch + [i]
for branch in acc
for i in arr
if sum(branch) < target
]
complete = [branch for branch in acc if sum(branch) == target]
return complete if not in_progress else go(in_progress + complete)
return go([[]])
Upvotes: 1
Reputation: 71451
You can use recursion with a generator:
def subsets(arr, _sum, c = []):
if sum(c) == _sum:
yield c
else:
for i in arr:
if sum(c+[i]) <= _sum:
yield from subsets(arr, _sum, c+[i])
print(list(subsets([1,3,5], 5)))
Output:
[[1, 1, 1, 1, 1], [1, 1, 3], [1, 3, 1], [3, 1, 1], [5]]
Upvotes: 1