sec.py

""" Some data about SEC football stadiums. """

from math import acos, cos, sin, pi

raw_data = [
    ("Florida",           "Gainesville, FL",  "Ben Hill Griffin Stadium",   29.6500583, -82.3487493),
    ("Georgia",           "Athens, GA",       "Sanford Stadium",            33.9497366, -83.3733819),
    ("Kentucky",          "Lexington, KY",    "Commonwealth Stadium",       38.0220905, -84.5053408),
    ("Missouri",          "Columbia, MO",     "Faurot Field",               38.9359174, -92.3334619),
    ("South Carolina",    "Columbia, SC",     "Williams-Brice Stadium",     33.9730840, -81.0187890),
    ("Tennessee",         "Knoxville, TN",    "Neyland Stadium",            35.9525345, -83.9246578),
    ("Vanderbilt",        "Nashville, TN",    "Dudley Field",               36.1430984, -86.8085619),
    ("Alabama",           "Tuscaloosa, AL",   "Bryant-Denny Stadium",       33.2079134, -87.5496290),
    ("Arkansas",          "Fayetteville, AR", "Reynolds Razorback Stadium", 36.0677677, -94.1778107),
    ("Auburn",            "Auburn, AL",       "Jordan-Hare Stadium",        32.6028163, -85.4901143),
    ("Louisiana State",   "Baton Rouge, LA",  "LSU Tiger Stadium",          30.4117210, -91.1842500),
    ("Mississippi State", "Starkville, MS",   "Scott Field",                33.4563214, -88.7938700),
    ("Mississippi",       "Oxford, MS",       "Vaught-Hemingway Stadium",   34.3618972, -89.5336923),
    ("Texas A&M",         "College Station",  "Kyle Field",                 30.6098417, -96.3404366)
]

names = [row[0] for row in raw_data]
coordinates_by_name = {uni: (lat, long) for uni, _, _, lat, long in raw_data}


def arclen(p1, p2):
    """ Compute the great circle distance, in radians, between two points
    p1 and p2, given as (latitude, longitude) pairs. """

    def toRadians(deg):
        return deg * pi / 180

    def toRect(p):
        lat, lng = p
        lat = toRadians(lat)
        lng = toRadians(lng)
        return (cos(lng) * cos(lat), sin(lng) * cos(lat), sin(lat))

    x1, y1, z1 = toRect(p1)
    x2, y2, z2 = toRect(p2)
    return acos(x1 * x2 + y1 * y2 + z1 * z2)

# Radius of the earth in miles.
earth_radius = 3956.6

def distance(u1, u2):
    return earth_radius * arclen(coordinates_by_name[u1], coordinates_by_name[u2])

# Actual driving distance is 345 miles; as the crow flies, 302.9.
assert 300 < distance("Florida", "Georgia") < 305

treeify.py

""" treeify.py - Find the minimum spanning tree of a graph. """

import heapq

def treeify(vertices, distanceFn):
    """ Return a list of triples (distance, v1, v2), the edges of a minimum spanning tree. """

    # This function implements Prim's algorithm:
    # https://en.wikipedia.org/wiki/Prim%27s_algorithm
    # We're going to build a tree. It is initially empty.
    tree_edges = []

    # `unreached` is the set of all vertices our tree hasn't reached yet, which
    # initially is all of them.
    unreached = set(vertices)

    # We'll begin by picking one vertex--it doesn't matter which one--and
    # putting it in our tree. This is the seed from which our tree will grow.
    seed = unreached.pop()

    # We use `edges` to select which edge to add next. Since the algorithm has
    # us repeatedly choosing the cheapest edge, we'll use a heap. Initially it
    # contains all edges leading out from `seed`.
    edges = []
    for v in unreached:
        heapq.heappush(edges, (distanceFn(seed, v), seed, v))

    # When no vertices are left unreached, we'll be done.
    while unreached:
        # Choose the cheapest edge remaining in `edges`.
        new_edge = d, x, y = heapq.heappop(edges)

        # `new_edge` might connect two already-reachable vertices; in that
        # case, skip it and try another. We know that `x` at least is in the
        # tree, since we've only ever added edges leading from inside the tree.
        assert x not in unreached
        if y in unreached:
            # Great, we have the cheapest edge that reaches a new vertex, as
            # required.  Add `y` and `new_edge` to the tree, and add to `edges`
            # all the edges leading from `y` to unreached vertices.
            tree_edges.append(new_edge)
            unreached.remove(y)
            for w in unreached:
                heapq.heappush(edges, (distanceFn(y, w), y, w))

    return tree_edges

if __name__ == '__main__':
    import sec, pprint
    result = treeify(sec.names, sec.distance)
    pprint.pprint(result)
    print "Total distance:", sum(row[0] for row in result)