2

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

  • all items assigned to bin 1, then bin 2, etc.
  • K-1 items assigned to bin 1, and one item to bin 2, etc.
  • ...

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.

Ribz
  • 551
  • 1
  • 4
  • 13

2 Answers2

6

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.

wim
  • 338,267
  • 99
  • 616
  • 750
  • I don't understand how this works. len(set(t)) would always return 2, no? Would you be able to explain what's going on and what you mean by the sort being 'stable' means? – freebie Nov 25 '16 at 23:17
  • @freebie As for the lambda function, `len(set(t))` will return 1 for `set((1,1))`, 2 for `set((1,2))`, etc. This is because in sets we do not have the same element more than once (unless it is a multiset). So, in fact, `set((1,1))={1}` in python. – Ribz Nov 25 '16 at 23:25
  • @Riebuoz ah yeah, I see. so the sort being 'stable' means that if elements are equivalent they will remain in their original location? Would it be much more complex to solve it otherwise? – freebie Nov 25 '16 at 23:31
  • Yes. You can see the explanation here also http://stackoverflow.com/a/1517824/3745674. I do not know about the complexity but I think that, in the worst case, it does not matter. – Ribz Nov 25 '16 at 23:38
  • Ah, so from that post it says the benefit is stacking for sorts. So if they started out unordered (`random.shuffle`) you could `L.sort()` then `L.sort(key=lambda t: len(set(t)))` – freebie Nov 25 '16 at 23:46
1

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.

jmd_dk
  • 12,125
  • 9
  • 63
  • 94