CODEWITHSUNDEEP
pythonalgorithmiterator
Lars
Lars

Reputation: 1969

algorithm to iterate through fixed size positive integer lists with the same sum

I am looking for a fast and memory efficient algorithm to iterate through all possible lists of positive integers of the same size (S) with a given sum (N).

for example if S = 3 and N = 4 the result would be (i have a really inefficient algo):

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

Not necessarily in that order. Another way to look at it is how to cut the number N into S pieces. The algorithm would be perfect if i could also set a maximum for each separate value in the lists.

I would use this to run through multidimensional arrays in a different order than produced by product(*indices).

Also generating all index combinations and sorting them by the sum would be too slow/memory consuming.

Upvotes: 0

Views: 759

Answers (2)

Lars
Lars

Reputation: 1969

Found a solution: it is based on the idea that a positive number N is a line of units and splitting them in S pieces is a matter of putting (S-1) separators in the list. These separators can iterated with combinations(range(N + S - 1), S - 1). The next step is to calculate the number of units before, between and after the separators:

def partition(N, size):
    n = N + size - 1
    for splits in combinations(range(n), size - 1):
        yield [s1 - s0 - 1 for s0, s1 in zip((-1,) + splits, splits + (n,))]

When you want to limit each item in the result you can filter unwanted out (surely not optimal, but i wanted to use combinations because it is probably implemented in C, and therefore probably much faster than anything i can come up with in python). The simpole version:

def sized_partition_slow(N, sizes):
    size = len(sizes)
    n = N + size - 1
    for splits in combinations(range(n), size - 1):
        result = [s1 - s0 - 1 for s0, s1 in zip((-1,) + splits, splits + (n,))]
        if all(r < s for r, s in zip(result, sizes)):
            yield result

And the faster but more complicated version:

def sized_partition(N, sizes):
    size = len(sizes)
    n = N + size - 1
    for splits in combinations(range(n), size - 1):
        result = []
        for s, s0, s1 in zip(sizes, (-1,) + splits, splits + (n,)):
            r = s1 - s0 - 1
            if r >= s:
                break
            result.append(r)
        else:
            yield result

I used this as early test: 

for indices in partition(4, 3):
    assert sum(indices) == 4
    assert all(0 <= i for i in indices)

for indices in sized_partition(4, [3, 3, 3]):
    assert sum(indices) == 4
    assert all(0 <= i < 3 for i in indices)

BTW: from the hip: you can generate the solution to the integer partitioning problem by iterating over S (size): as in:

def integer_partition(N, order=False):
    result = set()
    for size in range(1, N+1):
        for splits in combinations(range(1, N), size - 1):
            if order:
                p = tuple(s1 - s0 for s0, s1 in zip((0,) + splits, splits + (N,)))
            else:
                p = tuple(sorted(s1 - s0 for s0, s1 in zip((0,) + splits, splits + (N,))))
            result.add(p)
    return sorted(result, key=lambda r: (len(r), r))

I adapted the combinations() iterator a bit to not give zeros. It dedoubles for same partitions with different orders if order=False.

Upvotes: 2

Barış Can Tayiz
Barış Can Tayiz

Reputation: 87

This is simplest way I think:

def elect(S,N,List):
    result_list = []
    for list_val in List:
        if sum(list_val) == N:
            if len(list_val) == S:
                result_list.append(list_val)

    return result_list

this works under 1 sec for 1 million lists. If you want to fasten the way, you can put other if statements such as if sum(list_val[0:N/2]) > N or len(list_val) / 2 > S: such statements could detect faster the situtations.

Another way is sorting the lists and looking first N number sum. If it is greater than you want, you can elect these lists.

Upvotes: 0

Related Questions