brute.py

#!/usr/bin/env python

"""Bacon Graph Brute-Force Solver
By [email protected]

Synposis: python brute.py [filename]
    where filename is a valid bacon graph file

Reference: http://proggitquiz.com/challenge/4/
"""

import sys

# From mjard. Thanks.
def readmap(filename):
    cities = []
    fh = file(filename)

    rawdemensions, bacons = fh.readline().split(" ")
    rawmap = [list(l.strip()) for l in fh.readlines()]
    size_x, size_y = [int(x) for x in rawdemensions.split('x')]

    for y,row in enumerate(rawmap):
        for x,block in enumerate(row):
            if (block == 'P'):
                cities.append((x,y))

    fh.close()
    return (int(bacons), cities, (size_x, size_y))

def manhattan_distance(a, b):
    "Takes in two (x,y) tuples and returns the 'Manhattan' distance"
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def total_distance(point, city_list):
    "Finds the total distance between a given point and all cities"
    if point in city_list: raise Exception("Point is already a city!")
    return sum([manhattan_distance(point, xy) for xy in city_list])

def shortest_distance(point, dispenser_list):
    "Finds the distance between a given point and the closest dispenser"
    if point in dispenser_list: raise Exception("Point is on a dispenser!")
    return min([manhattan_distance(point, xy) for xy in dispenser_list])


def solve(num_dispensers, city_list, width, height):
    distances = []

    for y in xrange(height):
        for x in xrange(width):
            if not (x,y) in city_list:
                distances.append( ( total_distance((x,y), city_list), (x,y) ) )
    distances.sort(key=lambda x: x[0])

    # Strip off distances and just return coords
    return [d[1] for d in distances[:num_dispensers]]

def score(dispenser_list, city_list):
    "Returns sum of the distances from each city to the closest dispenser"
    return sum([shortest_distance(city, dispenser_list) for city in city_list])


def show_board(height, width, city_list, dispenser_list):
    for y in xrange(height):
        for x in xrange(width):
            if (x,y) in city_list:
                print 'P',
            elif (x,y) in dispenser_list:
                print '#',
            else: print '.',
        print

def main(infile):
    num_dispensers, city_list, (width, height) = readmap(infile)

    print "Loaded a valid %ix%i graph with %i cities" % (width, height, len(city_list))
    print "Adding %i bacon dispensers..." % num_dispensers

    dispenser_list = solve(num_dispensers, city_list, width, height)

    # Output the results
    print "Dispenser locations:", dispenser_list
    print "Score:", score(dispenser_list, city_list)

    show_board(height, width, city_list, dispenser_list)


if __name__ == '__main__':
    try: main(sys.argv[1])
    except: print __doc__