Riebuoz - 8 days ago 6

Python Question

I have

`K=2`

`N=3`

`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`

`N`

It is like having

`N`

`K`

- 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)`

`(2, 2)`

`(1, 2)`

Answer

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.

Source (Stackoverflow)

Comments