tsp.py

import sys
import matplotlib.pyplot as plt
import numpy as np
import scipy as sp
from scipy.spatial import distance_matrix
from scipy.optimize import linear_sum_assignment


sys.setrecursionlimit(1000)
BIG_M = 1e5


def detect_cycles(col_ind):

    """
    Returns
    -------
    List of arrays
    """
    n = len(col_ind)
    is_used = np.zeros(n, dtype=bool)
    cycles = []
    for i in range(n):
        if is_used[i]:
            continue
        source = i
        is_used[i] = True
        cycle = [i]
        while True:
            sink = col_ind[source]
            if is_used[sink]:
                break
            else:
                is_used[sink] = True
                cycle.append(sink)
                source = sink
        cycles.append(cycle)
    return cycles


def solve_tsp_exact(cost_matrix):
    # type: (np.ndarray) -> np.ndarray
    """
    Description
    -----------
    Solve Traveling Salesman Problem (TSP) and return exact solution.
    The algorithm is below:

        - Solve Assignment Problem to get relaxed solution of TSP
        - Use branch and bound method removing an edge which construct sub-cycle

    Parameters
    ----------
    cost_matrix:
        `distance_matrix[i, j]` is the distance between node i and j.

    Returns
    -------
    Return exact solution of TSP.
    solution:
        [4, 1, 3, 0, 2] means 4 -> 1 -> 3 -> 0 -> 2 -> 4 is the exact solution.
    """

    n = len(cost_matrix)
    cost_matrix[np.arange(n), np.arange(n)] = BIG_M
    current_best = np.inf
    solutions = []
    objective_values = []

    def solve_subproblem(cmat):
        # type: (np.ndarray) -> bool

        # if len(solutions) >= 10000:
        #     return True

        row_ind, col_ind = linear_sum_assignment(cmat)
        cycles = detect_cycles(col_ind)
        obj_val = cost_matrix[row_ind, col_ind].sum()

        # Terminate if there is no cycle and return the result
        if len(cycles) == 1:
            solution = np.array(cycles[0])
            solutions.append(solution)
            objective_values.append(obj_val)
            return True

        # Terminate if a deleted edge is used
        elif cmat[row_ind, col_ind].max() > BIG_M - 1e-2:
            return False

        # Terminate if the objective value of the relaxed problem is worse than the current best solution
        elif obj_val > current_best:
            return False

        # Solve subproblems
        else:
            # TODO: Choose the shortest cycle
            cycle = cycles[0]
            edges = list(zip(cycle, np.roll(cycle, -1)))
            # Create a copy of `cmat`
            cmat_ = np.array(cmat)

            for i, e in enumerate(edges):
                if i >= 1:
                    e_bef = edges[i-1]
                    cmat_[e_bef[0]] = BIG_M
                    cmat_[:, e_bef[1]] = BIG_M
                    cmat_[e_bef[0], e_bef[1]] = cmat[e_bef[0], e_bef[1]]
                # Remove an edge e
                cmat_[e[0], e[1]] = BIG_M
                solve_subproblem(cmat_)

    solve_subproblem(cost_matrix)
    print(len(solutions))
    return solutions[np.array(objective_values).argmin()]


# np.random.seed(42)
xs = sp.random.uniform(size=[8, 2])
cost_mat = distance_matrix(xs, xs)

route = solve_tsp_exact(cost_mat)

plt.plot(xs[:, 0], xs[:, 1], 'o')
plt.plot(xs[route][:, 0], xs[route][:, 1], '-')
plt.grid()
plt.show()