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?

Was it helpful?

Solution

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

OTHER TIPS

To find all assignments of N balls to M slots:

  • if N is 0 then
    • leave all M slots empty
  • otherwise, if M is 1, then
    • put all N balls to the only slot
  • otherwise
    • For each i in 0 .. N
      • put i balls in M-th slot, and
      • find all assignments of remaining N-i balls to remaining M-1 slots

Oh, here's a fun way to come at it:

>>> import random
>>> def permutate(N, M, lim = 1e6):
...     for i in range(int(lim)):
...         tup = ()
...         for j in range(M):
...             tup += (random.randrange(0,N+1),)
...
...         if sum(tup) == N: yield tup
...
>>> permutes = []
>>> for p in permutate(2, 3):
...     if not p in permutes:
...          permutes.append(p)
...
>>> permutes
[(0, 1, 1), (0, 0, 2), (2, 0, 0), (0, 2, 0), (1, 0, 1), (1, 1, 0)]
>>>

First we create a generator that gives a valid tuple of M elements which sums to N. We do this in a Monte Carlo-esk approach by composing each element of of a random number from 0 to N+1. We do this M times, and append each guess to a tuple. We then check if the tuple is valid (sums to N). We then call this generator repeatedly until it is exhausted, and add non-repeat entries to a list.

Some Things to Note:

  1. You'll need to tune lim to be appropriate for what you need. Generally it should be higher for larger M and N. You could modify this (with some effort) to estimate if lim is sized appropriate since you can figure how long permutes should be based off of N and M.

  2. Monte Carlo methods are slow, which is why I categorized this as a "fun" approach

Consider using itertools.combinations_with_replacement. You can easily do like this.

import itertools
def myfunc(n, m):
    for j in itertools.combinations_with_replacement(xrange(m), n):
        r = [0] * m
        for i in j:
            r[i] += 1
        yield r

The problem is how to ensure this is correct. I think @wnnmaw answer can be useful for that, because the logic is very simple.

#http://stackoverflow.com/a/22940260/2931409
import random
def permutate(N, M, lim = 1e6):
    for i in range(int(lim)):
        tup = []#<--Sorry changed here for test
        for j in range(M):
            tup += (random.randrange(0,N+1),)
        if sum(tup) == N: yield tup

def perm(n, m):
    permutes = []
    for p in permutate(n, m):
        if not p in permutes:
             permutes.append(p)
    return permutes

OK. Now check myfunc correctly works.

>>> N, M = 3, 5
>>> sorted(myfunc(N, M)) == sorted(perm(N, M))
True

Here's a nice way to do it.

NI = sum(in_pop)#number of input balls

# create the seed vector to be folded
u0 = [0] * NO
for i in range(NI):
    u0[i] = 1

# create a storage basis matrix and fold the seed vector
basis = [u0]
ui = []
for i in arange(1,NI):
    # copy the last vector
    ui = basis[i - 1][:]

    # find the number to 0 boundary
    edge = ui.index(0)

    # set the previous value to zero
    tmp = ui[edge - 1]
    ui[edge - 1] = 0

    # increment the value behind that
    ui[edge - 2] = ui[edge - 2] + tmp
    print 'edge @ ', edge
    print 'ui', ui

    basis.append(ui[:])

print basis

# run through the permutations in this basis
orderings = []
for ui in basis:
    perms = set(itertools.permutations(ui))
    orderings.extend(perms)

print orderings

The result is the following for input N=2, M=3:

edge @  2
ui [2, 0, 0]
[[1, 1, 0], [2, 0, 0]]
[(0, 1, 1), (1, 1, 0), (1, 0, 1), (0, 2, 0), (2, 0, 0), (0, 0, 2)]
[Finished in 0.5s]
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top