Suppose you have a list of the diameters of some circular holes, and a list of the diameters of some cylindrical pegs. Now the problem is finding all the ways you can fit the pegs in the holes, where the order in which you put them in doesn't matter, and it's OK to put a peg in a hole that's too large.

I've written the code below to solve the problem, and it works as intended. The innermost tuples of the answer matches a peg to a hole, i.e. (4, 7) means put d4 peg in d7 hole, and these tuples make up a set that represents one way of putting all the pegs in.

I'm not completely happy with the solution. It's not... snappy enough, and it feels overly complicated. Is there a smarter way to go about it?

def fit(pegs, holes):
    for h in holes:
        for p in pegs:
            if p <= h:
                pegs_left = [x for x in pegs if x != p]
                if pegs_left:
                    free_holes = [x for x in holes if x != h]
                    for rest in fit(pegs_left, free_holes):
                        yield frozenset(((p, h),)) | rest
                    yield frozenset(((p, h),))

def fitted(pegs, holes):
    return sorted(sorted(x) for x in set(fit(pegs, holes)))   

pegs = [2, 4, 7]
holes = [1, 3, 4, 6, 9]

print fitted(pegs, holes)
# [[(2, 3), (4, 4), (7, 9)], 
#  [(2, 3), (4, 6), (7, 9)], 
#  [(2, 4), (4, 6), (7, 9)], 
#  [(2, 6), (4, 4), (7, 9)]]
  1. I recommend against using single letter variable names, in particular your use of p and h. I think the code is clearer if you use peg and hole.
  2. Your code is going to break if you have multiple holes or pegs of the same size.
  3. You'd be better off using tuples instead of frozensets. You don't appear to be using the set features.

In order to have better code/algorithm you need to look at the problem differently. Right now you are trying to find peg/hole combinations. Instead, you should try to figure out where each peg fits.

You have the pegs: 2, 4, 7

Your goal is to select three of the holes: 1, 3, 4, 6, 9 which fit the pegs.

itertools.permutations(holes, 3)

will produce an iterator over all possible choices of 3 holes from the list. Something like:

1 3 4
1 3 6
1 3 9
1 4 3

We interpret each value as an assignment of holes to pegs. So 1,3,4 means (2, 1), (4, 3), (7,4). Of course not all such assignments are valid. So we filter out the invalid ones (i.e. where the pegs don't fit).

As an added bonus, itertools.permutations generates its values in sorted order. Because we only filter the results, the resulting function will also produce its values in sorted order. This removes the neccesity of doing the sorting yourself.

Here is my implementation:

import itertools

def check_fits(pegs, holes):
    for peg, hole in zip(pegs, holes):
        if peg > hole:
            return False
    return True

def fit(pegs, holes):
    assert len(pegs) <= len(holes)
    for selected_holes in itertools.permutations(holes, len(pegs)):
        if check_fits(pegs, selected_holes):
            yield zip(pegs, selected_holes)

pegs = [2, 4, 7]
holes = [1, 3, 4, 6, 9]

for item in fit(pegs, holes):
    print item

itertools.permutations will generate many invalid assignment of pegs to holes. For example, in the example above, it will generate a significant number of assignment which attempt to fit a peg of size 2 into a hole of size 1. A simple recursive function could avoid spending time on so many bad assignments. However, itertools.permutation is written in C, and so probably still has a speed advantage over such a recursive algorithm.

  • \$\begingroup\$ Thank you for a very nice answer! Timeit shows that your solution, in addition to being more readable and nicer overall, also is over an order of magnitude faster; 13.3 usec vs. 650 usec. A couple of comments: 1) Frozenset is there so the order in which the pegs are inserted won't matter, so it had a purpose! 2) In your solution, assert (...) should probably be replaced by if (...); raise ValueError(), since assert statements are ignored when python runs in optimized mode. \$\endgroup\$ – Lauritz V. Thaulow May 14 '11 at 17:47
  • \$\begingroup\$ Oh, and I suppose you meant "1,3,4 means (2, 1), (4, 3), (7, 4)"? \$\endgroup\$ – Lauritz V. Thaulow May 14 '11 at 17:53
  • \$\begingroup\$ @lazyr, re: frozenset, I understand now \$\endgroup\$ – Winston Ewert May 14 '11 at 18:47
  • \$\begingroup\$ @lazyr, re: numbers, you are correct. I've edited \$\endgroup\$ – Winston Ewert May 14 '11 at 18:47
  • \$\begingroup\$ @lazy, re: assert. it seems to me that depends on the context. Any sort of reusable function should raise a ValueError. However, if its strictly for internal use within one module then I'd probably use an assert. (In this case I was just trying to clarify that the code made that assumption. I didn't think too hard about it) \$\endgroup\$ – Winston Ewert May 14 '11 at 18:51

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.