all the ways to fit some pegs into given holes

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
                else:
                    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)]]

• 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.

• Your code is going to break if you have multiple holes or pegs of the same size.

• 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.

StackExchange Code Review Q#2407, answer score: 6