Travelling Salesman tour for ~33k cities with heuristic approximation - stanford algorithms thingy

tsp_heuristic.py

INF = 10 ** 9

class City(object):
    __slots__ = ('city_id', 'x', 'y')

    def __init__(self, city_id, x, y):
        self.city_id = int(city_id)
        self.x = x
        self.y = y

    def __repr__(self):
        return "{id}: ({x}, {y})".format(id=self.city_id, x=self.x, y=self.y)


def squared_distance(city1, city2):
    return ((city1.x - city2.x) ** 2) + ((city1.y - city2.y) ** 2)

def find_nearest_city(current_city, visited=None, all_cities=None):
    nearest = INF
    nearest_city = None
    for each_city in all_cities:
        if each_city.city_id not in visited and each_city != current_city:
            distance = squared_distance(current_city, each_city)
            if distance < nearest:
                nearest = distance
                nearest_city = each_city
    return (nearest, nearest_city)

if __name__ == '__main__':
    all_cities = []
    with open('nn.txt', 'r') as f:
        lines = f.readlines()
        n = int(lines[0])
        for line in lines[1:]:
            args = map(float, line.split(' '))
            all_cities.append(City(*args))

    source_city = all_cities[0]
    visited = {}
    ctr = 0
    total_distance = 0
    while ctr < n:
        path_distance, source_city = find_nearest_city(source_city, visited=visited, all_cities=all_cities)
        visited[source_city.city_id] = 1
        print "just visited city: ", source_city.city_id
        print "total visited cities: ", len(visited)
        total_distance += path_distance
        ctr += 1

    # go back to first city
    total_distance += squared_distance(all_cities[0], source_city)

    print int(total_distance ** (1 / 2.0))