recursive algorithm for balls in numbered boxes

Mark Dickinson mdickinson at enthought.com
Sun Sep 11 11:08:17 EDT 2011


On Sep 11, 1:43 am, "Dr. Phillip M. Feldman"
<Phillip.M.Feld... at gmail.com> wrote:
> I've written a recursive class that creates an iterator to solve a general
> formulation of the combinatorics problem known as "balls in numbered boxes"
> (also known as "indistinguishable balls in distinguishable boxes").  The
> code has been extensively tested and appears to work, but isn't terribly
> elegant.  Any suggestions about how to improve it will be appreciated.
>
> Also, I'd like to get this functionality into the Python's `itertools`
> module (the present set of combinatorics functions in `itertools` does not
> include "balls in boxes").  Does anyone know whom I should contact about
> this?

Note that for the version without size limits on individual boxes, the
itertools.combination function already provides most of what's
needed.  For example:

    import itertools

    def balls_in_boxes(nballs, nboxes):
        n, k = nballs + nboxes - 1, nboxes - 1
        for comb in itertools.combinations(range(n), k):
            yield [y - x - 1 for y, x in zip(comb + (n,), (-1,) +
comb)]

    print "5 balls in 3 boxes"
    for combination in balls_in_boxes(5, 3):
        print combination
        assert len(combination) == 3
        assert sum(combination) == 5


This is a well-known trick:  to divide 5 (unlabeled) balls amongst 3
(labeled) boxes, you write down sequences of 5 o's and 2 x's, where
the o's represent the 5 balls and the 'x's represent dividers:

    ooxooxo  -> [2, 2, 1]
    xoooxoo  -> [0, 3, 2]

And 'combinations(7, 2)' yields successively all the possible
different placements for the 2 dividers in the 7 symbols.


This question seems to come up often enough (without the box size
limit twist) that maybe it would be useful to include something like
this recipe in the itertool documentation.


For getting this into itertools, I'd suggest opening a feature request
on bugs.python.org and assigning it to Raymond Hettinger.

--
Mark



More information about the Python-list mailing list