MindSnacks engineering challenge 2

CirclePacker.py

#!/usr/bin/python -O
# Mike Solomon
# 2 Dec 2012
# MindSnacks engineering challenge 2

# usage: CirclePacker.py test_case_file

import math, copy, sys

class PackedCircle(object):
    ''' A circle to be packed in a tight rectangle '''
    def __init__(self, r, x=None):
        self.r = r
        self.x = x

    def getCenterRight(self, right):
        ''' Get center of circle packed snug against my right side '''
        return self.x + math.sqrt(4 * self.r * right.r)

    def checkIntersect(self, other):
        ''' Check if another circle intersects with me '''
        cent_distance_sq = (self.x - other.x)**2 + (self.r - other.r)**2
        # allow for a small amount of roundoff error
        return cent_distance_sq < (self.r + other.r)**2 - 1e-13


class CirclePacker(object):
    ''' Packs circles in a rectangle using simple backtracking '''
    def __init__(self):
        self.reset()

    def reset(self):
        ''' Remove best-so-far data from packer '''
        # the best-so-far box width and circle order
        self.best_width = float('inf')
        self.best_order = []

    def pack(self, radii, store_best=False):
        ''' Pack a set of circles and return size of box required '''
        self.reset()
        radii.sort(reverse=True)
        circles = map(PackedCircle, radii)
        self._packR(circles, store_best)
        return self.best_width

    def _packR(self, circles, store_best=False, width=0, order=[]):
        ''' Recursively pack circles '''
        if(len(circles) == 0): # done packing circles (basis case)
            if width < self.best_width:
                self.best_width = width
                if store_best:
                    # make a deep copy, or we'll lose the x values
                    self.best_order = copy.deepcopy(order)
                return

        # remember sizes we've already fitted so we can prune
        sizes_tried = set()
        # pack the rest of the circles (general case)
        for circle_number, circle in enumerate(circles):
            if circle.r in sizes_tried:
                # skip, we've already tried an equal sized circle
                continue
            sizes_tried.add(circle.r)
            packed = False

            if(len(order) == 0): # first circle
                circle.x = circle.r
                packed = True
            else: # a subsequent circle
                # try placing adjacent and right of every circle
                for other in reversed(order):
                    circle.x = other.getCenterRight(circle)
                    # check if we cross the y axis
                    left_edge = circle.x - circle.r
                    if left_edge < 0: continue
                    # check if we overlap when packed here
                    packed = True

                    for c in order:
                        # do fast left edge check, then slower accurate check
                        if left_edge < c.right_edge and c.checkIntersect(circle):
                            # bail as soon as we know we intersect
                            packed = False 
                            break
                    if packed: break # we managed to place ourself
            if not packed:
                # impossible to place this circle given the rest
                continue
            circle.right_edge = circle.x + circle.r # stored for fast edge check
            # width is either same as before, or extended by new circle
            width = max(width, circle.right_edge)
            if width < self.best_width: # could still be best
                # continue the search
                self._packR(circles[:circle_number] + circles[circle_number+1:],
                            store_best,
                            width,
                            order + [circle])
                

def readCircleFile(fname):
    ''' Yield test cases from a file one at a time '''
    with open(fname, 'r') as f:
        for line in f:
            # break into radii tokens
            str_radii = line.split()[1:]
            # skip non- test cases
            if(len(str_radii) > 0):
                # convert strings to floats
                yield map(float, str_radii)


def packAndPrintFile(fname, plot):
    ''' Pack and print results of circle packing, optionally plotting ''' 
    cp = CirclePacker()
    if plot:
        import matplotlib.pyplot as plt

    for c_num, case in enumerate(readCircleFile(fname)):
        width = cp.pack(case, plot)
        print '{:.3f}'.format(width)
        if plot:
            fig = plt.gcf()
            ax = plt.gca()
            ax.set_ylim(0, max(case) * 2)
            ax.set_xlim(0, cp.best_width)
            ax.set_aspect('equal')
            for i, circle in enumerate(cp.best_order):
                c = plt.Circle((circle.x, circle.r), circle.r)
                c.set_alpha(0.2)
                fig.gca().add_artist(c)
                #ax.axvline(circle.x + circle.r)
            fig.savefig('{}.png'.format(c_num))
            fig.clf()

        
if __name__ == '__main__':
    plot = False # requires matplotlib
    if len(sys.argv) < 2:
        print >> sys.stderr, 'usage: CirclePacker.py test_case_file'
    else:
        # skip most of the usual error checks, for brevity
        packAndPrintFile(sys.argv[1], plot)