Create a gist now

Instantly share code, notes, and snippets.

@jorendorff /sec.py
Last active Dec 14, 2015

What would you like to do?
""" 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 - 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment