Travelling salesman using brute-force and heuristics

Question

I have implemented both a brute-force and a heuristic algorithm to solve the travelling salesman problem.

import doctest
from itertools import permutations


def distance(point1, point2):
    """
    Returns the Euclidean distance of two points in the Cartesian Plane.

    >>> distance([3,4],[0,0])
    5.0
    >>> distance([3,6],[10,6])
    7.0
    """
    return ((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2) ** 0.5


def total_distance(points):
    """
    Returns the length of the path passing throught
    all the points in the given order.

    >>> total_distance([[1,2],[4,6]])
    5.0
    >>> total_distance([[3,6],[7,6],[12,6]])
    9.0
    """
    return sum([distance(point, points[index + 1]) for index, point in enumerate(points[:-1])])


def travelling_salesman(points, start=None):
    """
    Finds the shortest route to visit all the cities by bruteforce.
    Time complexity is O(N!), so never use on long lists.

    >>> travelling_salesman([[0,0],[10,0],[6,0]])
    ([0, 0], [6, 0], [10, 0])
    >>> travelling_salesman([[0,0],[6,0],[2,3],[3,7],[0.5,9],[3,5],[9,1]])
    ([0, 0], [6, 0], [9, 1], [2, 3], [3, 5], [3, 7], [0.5, 9])
    """
    if start is None:
        start = points[0]
    return min([perm for perm in permutations(points) if perm[0] == start], key=total_distance)


def optimized_travelling_salesman(points, start=None):
    """
    As solving the problem in the brute force way is too slow,
    this function implements a simple heuristic: always
    go to the nearest city.

    Even if this algoritmh is extremely simple, it works pretty well
    giving a solution only about 25% longer than the optimal one (cit. Wikipedia),
    and runs very fast in O(N^2) time complexity.

    >>> optimized_travelling_salesman([[i,j] for i in range(5) for j in range(5)])
    [[0, 0], [0, 1], [0, 2], [0, 3], [0, 4], [1, 4], [1, 3], [1, 2], [1, 1], [1, 0], [2, 0], [2, 1], [2, 2], [2, 3], [2, 4], [3, 4], [3, 3], [3, 2], [3, 1], [3, 0], [4, 0], [4, 1], [4, 2], [4, 3], [4, 4]]
    >>> optimized_travelling_salesman([[0,0],[10,0],[6,0]])
    [[0, 0], [6, 0], [10, 0]]
    """
    if start is None:
        start = points[0]
    must_visit = points
    path = [start]
    must_visit.remove(start)
    while must_visit:
        nearest = min(must_visit, key=lambda x: distance(path[-1], x))
        path.append(nearest)
        must_visit.remove(nearest)
    return path


def main():
    doctest.testmod()
    points = [[0, 0], [1, 5.7], [2, 3], [3, 7],
              [0.5, 9], [3, 5], [9, 1], [10, 5]]
    print("""The minimum distance to visit all the following points: {}
starting at {} is {}.

The optimized algoritmh yields a path long {}.""".format(
        tuple(points),
        points[0],
        total_distance(travelling_salesman(points)),
        total_distance(optimized_travelling_salesman(points))))


if __name__ == "__main__":
    main()

Answer 1

I enjoyed the first look at the code as it's very clean, you have extensive docstrings and great, expressive function names. Now you know the deal with PEP8, but except for the one 200 character long line I don't think it matters much really.

There're a few typo with the wrong spelling "algoritmh".

The coordinates should be immutable 2-tuples. The reason being the safety of immutable data-structures. YMMV, but that makes it really obvious that those are coordinates as well.

optimized_travelling_salesman should make a defensive copy of points, or you should otherwise indicate that it's destructive on that argument.

Instead of if start is None: start = points[0] you could also use start = start or points[0] to save some space while still being relatively readable.

For the algorithms the only thing I'd is not to use square root if you don't have to. You can basically create a distance_squared and use that instead of distance because the relationship between a smaller and bigger distance will stay the same regardless. That doesn't apply for the final output of course. Edit: And, as mentioned below by @JanneKarila, you can't use that for the brute-force version.

Answer 2

Besides thanking you for posting this, I wanted to share how I edited your code to do something I needed.

I'm implementing a survey for econ research. Part of the implementation is giving our survey team a list of addresses so that they can visit them. After geo-coding these addresses, we organized them into neighborhoods through a very simple k-means clustering process. Each member of our team is assigned a zone, therefore we needed to provide them with an optimized route within each zone.

The issue at hand was that members of our survey need additional data on top of the location data - such as unique IDs for households and basic contact info for the household. So, somehow I needed to keep some form of identification for each coordinate so I could merge it back with the relevant data.

I realize this might be a simple edit, but I couldn't figure out another way to have a relation between the two tables of data.

The code is exactly the same, except for that it returns an enumerated list (enumerated at the time of input) so that you can merge it back with other data.

I also added a haversine formula for calculating the distance between two geo-coordinates in miles.

import numpy as np
def haversine_enum(item1, item2):
    """
    Returns the great-circle distance between two enumearted points
    on a sphere given their indexes, longitudes, and latitudes in the
    form of a tuple.

    >>> haversine_enum((0, (3,5)), (1, (4,7)))
    154.27478490048566

    >>> haversine_enum((0, (41.325, -72.325)), (1, (41.327, -72.327)))
    0.17282397386672291
    """
    r = 3959 #radius of the earth
    r_lat = np.radians(item1[1][0])
    r_lat2 = np.radians(item2[1][0])
    delta_r_lat = np.radians(item2[1][0]-item1[1][0])
    delta_r_lon = np.radians(item2[1][1]-item1[1][1])

    a = (np.sin(delta_r_lat / 2) * np.sin(delta_r_lat / 2) +
        np.cos(r_lat) * np.cos(r_lat2) *
        np.sin(delta_r_lon / 2) * np.sin(delta_r_lon / 2))

    c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1-a))

    d = r * c

    return d

def optimized_travelling_salesman_enum(points, start=None):
    """
    Taken from:
        https://codereview.stackexchange.com/questions/81865/
        travelling-salesman-using-brute-force-and-heuristics

    As solving the problem in the brute force way is too slow,
    this function implements a simple heuristic: always
    go to the nearest city.

    Even if this algoritmh is extremely simple, it works pretty well
    giving a solution only about 25% longer than the optimal one (cit. Wikipedia),
    and runs very fast in O(N^2) time complexity.
    """

    points = list(enumerate(points))
    if start is None:
        start = points[0]

    must_visit = points
    path = [start]
    must_visit.remove(start)

    while must_visit:
        nearest = min(must_visit,
                       key=lambda x: haversine_enum(path[-1], x))
        path.append(nearest)
        must_visit.remove(nearest)

    return path