Every way to organize N objects in M list slots

Question

I'm trying to find a way, using built-in functions, to list every way to organize N balls in M slots. The balls can stack in the slots. For example:

N = 2, M = 3 -> {|0|1|1|, |1|0|1|, |1|1|0|, |2|0|0|, |0|2|0|, |0|0|2|}

itertools.permutations() is part of the puzzle, but how can you go through all possible stacks of balls that preserves N?

Answer 1

Let oo represent two balls. The problem of placing 2 balls in 3 slots is equivalent to the problem of placing 2 balls around 2 sticks:

|o|o   -->  0,1,1
o||o        1,0,1
o|o|        1,1,0
oo||        2,0,0
|oo|        0,2,0
||oo        0,0,2

The sticks, | represent the edges of the slots. The number of balls in each slot is shown on the right.

Notice that in each row there are 4 locations, and 2 sticks. There is always one fewer stick than there are slots. So the problem is equivalent to finding all the ways to select 2 locations for the balls out of 4 possibilities. 4 choose 2.

In [98]: import itertools as IT

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

These, then, are the possible locations for the balls. All that's left to do is to compute into which slot these balls belong. Let's take (1, 3) as an example. The pictoral diagram for (1, 3) looks like this:

|o|o

It turns out that if you subtract (0, 1) elementwise from (1, 3), you get the slot for each ball:

(1, 3)
(0, 1)
------
(1, 2)

Thus, the first ball is in slot 1, the second in slot 2. In general, if you subtract range(m-1) from the combination, you get the slot values. This makes some sense if you think of the amount you are subtracting as the number of balls that precede the current ball in the pictoral diagram. Since our diagrams consist of balls and slots, if you subtract the balls, what remains are slots.

---

import itertools as IT
def pigeonhole(n, m):
    """
    Generate the ways n balls can placed in m slots
    """
    for choice in IT.combinations(range(n+m-1), n):
        slot = [c-i for i,c in enumerate(choice)]
        result = [0]*m
        for i in slot:
            result[i] += 1
        yield result            

print(list(pigeonhole(2,3)))

yields

[[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0], [0, 1, 1], [0, 0, 2]]