MyGGaN
MyGGaN

Reputation: 1806

set-wise combinations in Python3

I have four sets (a-d in this example) and I'd like to iterate through all combinations with the restriction that you can only pick one element from each set and you need to pick one from all sets. In the case of four sets you could simply do something like this:

for a in (1,2):
    for b in (3,4):
        for c in (5,6):
            for d in (7,8):
                print(a,b,c,d)

And the expected output would be:

1 3 5 7
1 3 5 8
1 3 6 7
1 3 6 8
1 4 5 7
1 4 5 8
1 4 6 7
1 4 6 8
2 3 5 7
2 3 5 8
2 3 6 7
2 3 6 8
2 4 5 7
2 4 5 8
2 4 6 7
2 4 6 8

How could I rewrite this for n-sets?

I'm hoping for something like itertools.combinations(). Or iterating over a n-dimensional array. Probably nympy have something?

Thanks

Upvotes: 1

Views: 61

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53029

Here is a numpy solution:

a = tuple(np.arange(x,x+2) for x in range(0,8,2))
a
# (array([0, 1]), array([2, 3]), array([4, 5]), array([6, 7]))
for x,y,z,u in np.broadcast(*np.ix_(*a)):
    print(x,y,z,u) 
# 0 2 4 6
# 0 2 4 7
# 0 2 5 6
# 0 2 5 7
# 0 3 4 6
...

Upvotes: 1

MyGGaN
MyGGaN

Reputation: 1806

As Paul Panzer suggested in a comment. This is itertools.product() and the example would look like this:

import itertools
for p in itertools.product((1,2), (3,4), (5,6), (7,8)):
    print(p)

(1, 3, 5, 7)
(1, 3, 5, 8)
(1, 3, 6, 7)
(1, 3, 6, 8)
(1, 4, 5, 7)
(1, 4, 5, 8)
(1, 4, 6, 7)
(1, 4, 6, 8)
(2, 3, 5, 7)
(2, 3, 5, 8)
(2, 3, 6, 7)
(2, 3, 6, 8)
(2, 4, 5, 7)
(2, 4, 5, 8)
(2, 4, 6, 7)
(2, 4, 6, 8)

Upvotes: 0

Related Questions