CODEWITHSUNDEEP
pythoncombinations
Ribz
Ribz

Reputation: 551

How to sort all combinations in python?

I have K=2 and N=3 and I generate all combinations as follows:

list(itertools.product(range(1, N+1), repeat=K))

and I get

[(1, 1), 
 (1, 2),
 (1, 3), 
 (2, 1), 
 (2, 2), 
 (2, 3), 
 (3, 1), 
 (3, 2), 
 (3, 3)]

I need to sort these combinations to get 

[(1, 1), 
 (2, 2),
 (3, 3),
 (1, 2),
 (1, 3), 
 (2, 1), 
 (2, 3), 
 (3, 1), 
 (3, 2)]

How can I do this for general K and N?

It is like having N bins and K items and I would like to produce all possible assignment of items to bins but starting with

So in the example (1, 1) means that all items are in bin 1, (2, 2) means that all items are in bin 2, etc. (1, 2) means that item 1 is in bin 1 and item 2 is in bin 2, etc.

Upvotes: 2

Views: 1368

Answers (2)

jmd_dk
jmd_dk

Reputation: 13120

You can simply do this:

import itertools
K = 2
N = 3
a = list(itertools.product(range(1, N+1), repeat=K))
a_sorted = sorted(a, key=lambda t: -(t[0] == t[1]))

The lambda function tells sorted to place the tuples with equal numbers in the beginning. We also need to sort after the numeric value, but as this is the default behavior of sorted, we do not need to add this in by hand.

Upvotes: 1

wim
wim

Reputation: 363233

It is already generated almost how you wanted, so you can take advantage of python's sort being stable:

>>> L = list(itertools.product(range(1, N+1), repeat=K))
>>> L.sort(key=lambda t: len(set(t)))
>>> L
[(1, 1), (2, 2), (3, 3), (1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

This just pushes the tuples with the most equal values towards the front. It should generalise to higher dimensions consistent with the way you've described.

Upvotes: 6

Related Questions