greedy.py

from __future__ import print_function, division
from itertools import islice
from array import array as pyarray
################################################################################
# A simple algorithm for solving the Travelling Salesman Problem
# Finds a suboptimal solution
################################################################################
if "xrange" not in globals():
    #py3
    xrange = range
else:
    #py2
    pass
    

def optimize_solution( distances, connections, endpoints ):
    """Tries to optimize solution, found by the greedy algorithm"""
    N = len(connections)
    path = restore_path( connections, endpoints )
    def ds(i,j): #distance between ith and jth points of path
        pi = path[i]
        pj = path[j]
        if pi < pj:
            return distances[pj][pi]
        else:
            return distances[pi][pj]
            
    d_total = 0.0
    optimizations = 0
    for a in xrange(N-1):
        b = a+1
        for c in xrange( b+2, N-1):
            d = c+1
            delta_d = ds(a,b)+ds(c,d) -( ds(a,c)+ds(b,d))
            if delta_d > 0:
                d_total += delta_d
                optimizations += 1
                connections[path[a]].remove(path[b])
                connections[path[a]].append(path[c])
                connections[path[b]].remove(path[a])
                connections[path[b]].append(path[d])

                connections[path[c]].remove(path[d])
                connections[path[c]].append(path[a])
                connections[path[d]].remove(path[c])
                connections[path[d]].append(path[b])
                path[:] = restore_path( connections, endpoints )
    
    return optimizations, d_total
        
def restore_path( connections, endpoints ):
    """Takes array of connections and returns a path.
    Connections is array of lists with 1 or 2 elements.
    These elements are indices of teh vertices, connected to this vertex
    Guarantees that first index < last index
    """
    #when endpoints is None, then both start and end are not specified.
    start, end = endpoints or (None, None)
    #Now, if start, end or both are not specified - replace them with  replace unspecified start or end endpoints with the found ones
    need_revert = False
    is_loop = (start is not None) and (start == end)
    if start is None:
        if end is None:
            #find first node that only have one link.
            start = next(idx 
                         for idx, conn in enumerate(connections)
                         if len(conn)==1)
        else:
            #In this case - search from the end, then reverse the order
            start = end
            need_revert = True
    
    #ready to generate the path now.
    path = [start]
    prev_point = None
    cur_point = start
    #We know that path len should be the same as number of connection nodes, or one more if we are searching for a loop.
    #We already have one node, so need the rest
    for _ in xrange(len(connections) - (0 if is_loop else 1)):
        next_point = next(pnt for pnt in connections[cur_point] 
                          if pnt != prev_point )
        path.append(next_point)
        prev_point, cur_point = cur_point, next_point
    if need_revert:
        return path[::-1]
    else:
        return path

def _assert_triangular(distances):
    """Ensure that matrix is left-triangular at least.
    """
    for i, row in enumerate(distances):
        if len(row) < i: raise ValueError( "Distance matrix must be left-triangular at least. Row {row} must have at least {i} items".format(**locals()))
    

def pairs_by_dist(N, distances):
    """returns list of coordinate pairs (i,j), sorted by distances; such that i < j"""
    #Sort coordinate pairs by distance
    indices = []
    for i in xrange(N):
        for j in xrange(i):
            indices.append(i*N+j)

    indices.sort(key = lambda ij: distances[ij//N][ij%N])
    return ((ij//N,ij%N) for ij in indices)

def solve_tsp( distances, optim_steps=3, pairs_by_dist=pairs_by_dist, endpoints=None ):
    """Given a distance matrix, finds a solution for the TSP problem.
    Returns list of vertex indices. 
    Guarantees that the first index is lower than the last

    :arg: distances : left-triangular matrix of distances. array of arrays
    :arg: optim_steps (int) number of additional optimization steps, allows to improve solution but costly.
    :arg: pairs_by_dist (function) an implementtion of the pairs_by_dist function. for optimization purposes.
    :arg: endpoinds : None or pair (int or None, int or None). Specifies start and end nodes of the path. None is unspecified.
    """
    N = len(distances)
    start, end = endpoints or (None, None)
    #When both points are specified, we are looking for a loop.
    is_loop = (start is not None) and (start == end)
    if start is not None and not (0<=start<N): raise ValueError("Start point does not belong range")
    if end is not None and not (0<=end<N): raise ValueError("Start point does not belong range")
        
    if N == 0: return []
    if N == 1:
        return [0,0] if is_loop else [0]
    if N == 2 and is_loop:
        #Too short loops are easier to bypass here
        return [start, 1-start, start]
    _assert_triangular(distances)

    #State of the TSP solver algorithm.
    node_valency = pyarray('i', [2])*N #Initially, each node has 2 sticky ends
    has_both_endpoints = (start is not None) and (end is not None)

    if not is_loop:
        #in a loop, all nodes have valency 2. Otherwise - start and end has 1
        if start is not None:
            node_valency[start]=1
        if end is not None:
            node_valency[end]=1
        
    #for each node, stores 1 or 2 connected nodes
    connections = [[] for i in xrange(N)] 

    def join_segments(sorted_pairs):
        #segments of nodes. Initially, each segment contains only 1 node
        segments = [ [i] for i in xrange(N) ]
  
        def possible_edges():
            #Generate sequence of graph edges, that are possible and connect different segments.
            for ij in sorted_pairs:
                i,j = ij
                #if both start and end could have connections,
                #  and both nodes connect to a different segments:
                if node_valency[i] and node_valency[j] and\
                   (segments[i] is not segments[j]): 
                    yield ij
                    
        def connect_vertices(i,j):
            node_valency[i] -= 1
            node_valency[j] -= 1
            connections[i].append(j)
            connections[j].append(i)
            #Merge segment J into segment I.
            seg_i = segments[i]
            seg_j = segments[j]
            if len(seg_j) > len(seg_i):
                seg_i, seg_j = seg_j, seg_i
                i, j = j, i
            for node_idx in seg_j:
                segments[node_idx] = seg_i
            seg_i.extend(seg_j)
            
        def edge_connects_endpoint_segments(i,j):
            #return True, if given ede merges 2 segments that have endpoints in them
            #when this happens, search would terminate prematurely.
            #ALso works with the case when both endpoints are the same.
            si,sj = segments[i],segments[j]
            ss,se = segments[start], segments[end]
            return (si is ss) and (sj is se) or (sj is ss) and (si is se)
                
            
        #Take first N-1 possible edge. they are already sorted by distance
        edges_left = N-1
        for i,j in possible_edges():
            if has_both_endpoints and edges_left!=1 and edge_connects_endpoint_segments(i,j):
                continue #don't allow premature path termination
            connect_vertices(i,j)
            edges_left -= 1
            if edges_left == 0:
                break
        #if searching for a loop then close it
        if is_loop:
            _close_loop(connections)
    #invoke main greedy algorithm
    join_segments(pairs_by_dist(N, distances))

    #now call additional optiomization procedure.
    for passn in range(optim_steps):
        nopt, dtotal = optimize_solution( distances, connections, endpoints )
        if nopt == 0:
            break
    #restore path from the connections map (graph) and return it
    return restore_path( connections, endpoints=endpoints )

def _close_loop(connections):
    """Modify connections to close the loop"""
    i,j = (i for i, conn in enumerate(connections)
           if len(conn)==1)
    connections[i].append(j)
    connections[j].append(i)