CODEWITHSUNDEEP
pythonrecursiongenerator
Jonathan
Jonathan

Reputation: 1088

Rewriting a function that creates combinations of numbers with a fixed sum as a generator

I want to turn a recursive function into a generator in python. Currently, I have this function to create combinations of numbers with a fixed sum,

def combinations_fixed_sum(fixed_sum, length_of_list, lst=[]):
  if length_of_list == 1:
    lst += [fixed_sum]
    print(lst)
  else:
    for i in range(fixed_sum+1):
      combinations_fixed_sum(i, length_of_list-1, lst + [fixed_sum-i])

The print statement is the part, that I want to return to the generator. Is this possible?

Upvotes: 3

Views: 243

Answers (3)

Nelewout
Nelewout

Reputation: 6574

This is not too difficult to achieve, once we look at the problem a little differently. The first thing to notice is that once n - 1 numbers have been determined, the last number is fixed also, as all numbers together need to sum to some constant. This implies our search space is actually in n - 1 dimensions, rather than the full n.

Why does this matter, you say? Well, this implies we can proceed with normal iteration over this search space, without having to take into account that the numbers must sum to some constant: we can determine the final number ad hoc!

I'll first present the code, and then explain how it works,

import itertools
from collections import deque


def neighbours(node):       # see https://stackoverflow.com/a/45618158/4316405
    for relative_index in itertools.product((0, 1), repeat=len(node)):
        yield tuple(i + i_rel for i, i_rel
                    in zip(node, relative_index))


def combinations(total, numbers):
    nodes = deque([tuple([0] * (numbers - 1))])  # origin
    seen = set()

    while len(nodes) != 0:                       # while not empty
        node = nodes.popleft()

        for neighbour in neighbours(node):
            if neighbour in seen:
                continue

            if sum(neighbour) > total:
                continue                         # should not go out-of-bounds

            nodes.append(neighbour)

            seen.add(neighbour)

            yield [*neighbour, total - sum(neighbour)]

Here we walk along an (n - 1)-dimensional grid, starting at the origin and walking around to visit neighbouring grid points (keeping track of those we have seen before). These new grid points are then visited, and the final entry is determined as well. We keep adding grid points, so long as their indices do not sum to a value larger than total. Once we have exhausted all such admissible grid points, we are done iterating.

Upvotes: 0

kuppern87
kuppern87

Reputation: 1135

It's possible. The most straight forward way is to replace the print with a yield and then make sure to iterate through the generators returned by all the calls to the recursive function. 

def combinations_fixed_sum_chain(fixed_sum, length_of_list, lst=[]):
    if length_of_list == 1:
        lst += [fixed_sum]
        yield lst # 1 element generator
    else:
        for i in range(fixed_sum+1):
            # Iterate through each generator
            for comb in combinations_fixed_sum(i, length_of_list-1, lst + [fixed_sum-i]):
                yield comb

Or you could use itertools.chain.from_iterable to chain all the generators together:

from itertools import chain
def combinations_fixed_sum(fixed_sum, length_of_list, lst=[]):
    if length_of_list == 1:
        lst += [fixed_sum]
        return (lst,)
    return chain.from_iterable(
        combinations_fixed_sum(i, length_of_list-1, lst + [fixed_sum-i])
        for i in range(fixed_sum+1))

Upvotes: 0

niemmi
niemmi

Reputation: 17263

You can use yield in return individual result to caller. Recursion can be handled with yield from that will yield all the values from nested generator:

def combinations_fixed_sum(fixed_sum, length_of_list, lst=[]):
  if length_of_list == 1:
    lst += [fixed_sum]
    yield lst
  else:
    for i in range(fixed_sum+1):
      yield from combinations_fixed_sum(i, length_of_list-1, lst + [fixed_sum-i])

print(list(combinations_fixed_sum(4, 2)))

Output:

[[4, 0], [3, 1], [2, 2], [1, 3], [0, 4]]

Note that yield from is only available on Python 3, if you're using Python 2.x you need to yield values individually:

def combinations_fixed_sum(fixed_sum, length_of_list, lst=[]):
  if length_of_list == 1:
    lst += [fixed_sum]
    yield lst
  else:
    for i in range(fixed_sum+1):
      for x in combinations_fixed_sum(i, length_of_list-1, lst + [fixed_sum-i]):
        yield x

Upvotes: 1

Related Questions