CODEWITHSUNDEEP
pythoncombinations
user2297541
user2297541

Reputation: 13

Python list of combinations of positions

I'm trying to generate a list of all possible 1-dimensional positions for an arbitrary number of identical objects. I want it formatted so each coordinate is the distance from the previous object, so for 3 objects (0,5,2) would mean one object is at position 0, another is at position 5 and another is at position 7.

So the main restraint is that the sum of the coordinates is <=D. Nested for loops works well for this. For example, with 3 objects with maximum coordinate D:

def positions(D):
    output=[]
    for i in range(D+1):
        for j in range(D+1-i):
            for k in range(D+1-i-j):
                 output.append((i,j,k))
     return(output)

What's the best way to extend this to an arbitrary number of objects? I can't find a good way without explicitly writing a specific number of for loops.

Upvotes: 1

Views: 911

Answers (1)

DSM
DSM

Reputation: 353419

I think you can combine itertools.combinations, which will give you the locations, with taking the difference, which should give you your "distance from the previous object" behaviour. For example, using

def diff(loc):
    return [y-x for x,y in zip((0,) + loc, loc)]

we have

In [114]: list(itertools.combinations(range(4), 3))
Out[114]: [(0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)]

for the possible positions, and then

In [115]: [diff(x) for x in itertools.combinations(range(4), 3)]
Out[115]: [[0, 1, 1], [0, 1, 2], [0, 2, 1], [1, 1, 1]]

for your relative-distance version.

Upvotes: 3

Related Questions