gal leshem
gal leshem

Reputation: 561

Recursive program to get all subsets with given sum includes repetitions

I am trying to write a program that takes as input:

  1. a list of allowed numbers (arr)
  2. a total (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

Answers (2)

joel
joel

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

Ajax1234
Ajax1234

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

Related Questions