CSE505/Assignment2/A/tsp.py
2024-03-09 23:16:18 +08:00

105 lines
4.2 KiB
Python

"""Simple Travelling Salesperson Problem (TSP) between cities."""
from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp
import timeit
def create_data_model():
"""Stores the data for the problem."""
data = {}
data["distance_matrix"] = [
[8, 8, 15, 7, 11, 17, 8, 3, 4, 11, 9, 6, 2, 7, 15, 6, 16, 5, 9, 14],
[16, 17, 13, 16, 15, 17, 12, 11, 7, 8, 13, 3, 7, 15, 9, 8, 16, 18, 7, 5],
[16, 4, 1, 9, 17, 11, 18, 2, 16, 13, 6, 5, 2, 7, 18, 16, 19, 14, 14, 9],
[9, 9, 7, 12, 16, 12, 3, 5, 16, 18, 10, 7, 6, 12, 4, 5, 14, 11, 5, 5],
[5, 2, 16, 18, 15, 15, 8, 10, 15, 18, 10, 7, 2, 1, 15, 17, 2, 15, 9, 6],
[9, 7, 11, 10, 12, 6, 10, 1, 5, 20, 19, 12, 13, 1, 19, 16, 17, 20, 16, 2],
[15, 7, 3, 15, 20, 5, 10, 4, 20, 13, 5, 11, 20, 16, 13, 4, 7, 11, 13, 17],
[11, 16, 19, 15, 12, 12, 10, 9, 1, 15, 17, 17, 18, 17, 9, 12, 9, 4, 17, 3],
[2, 5, 16, 14, 5, 3, 16, 7, 3, 2, 1, 3, 20, 16, 14, 9, 10, 18, 20, 8],
[9, 4, 20, 2, 11, 19, 5, 3, 11, 2, 13, 20, 17, 5, 12, 17, 17, 9, 5, 4],
[5, 14, 7, 14, 3, 13, 1, 13, 17, 6, 18, 18, 16, 15, 18, 19, 15, 12, 9, 4],
[14, 2, 17, 8, 11, 16, 7, 12, 12, 6, 13, 12, 6, 18, 3, 7, 3, 1, 7, 17],
[8, 10, 8, 13, 15, 20, 2, 11, 9, 1, 4, 14, 1, 19, 15, 13, 18, 16, 10, 2],
[16, 2, 14, 3, 11, 17, 7, 6, 4, 2, 3, 8, 20, 8, 7, 12, 1, 4, 20, 10],
[5, 19, 10, 18, 6, 18, 2, 11, 6, 18, 11, 18, 13, 7, 18, 15, 3, 10, 20, 7],
[16, 19, 12, 19, 15, 12, 14, 8, 7, 11, 4, 13, 19, 12, 14, 13, 3, 7, 8, 20],
[4, 4, 15, 5, 16, 14, 18, 11, 18, 20, 19, 18, 18, 4, 3, 20, 5, 18, 18, 1],
[19, 6, 2, 15, 20, 3, 13, 7, 10, 17, 12, 11, 10, 19, 6, 2, 2, 19, 15, 1],
[16, 16, 3, 7, 6, 12, 6, 3, 2, 1, 11, 8, 13, 16, 14, 5, 20, 20, 10, 3],
[3, 14, 13, 20, 20, 2, 7, 14, 2, 6, 2, 19, 20, 20, 14, 7, 7, 5, 15, 18],
]
data["num_vehicles"] = 1
data["depot"] = 0
return data
def print_solution(manager, routing, solution):
"""Prints solution on console."""
print(f"Objective: {solution.ObjectiveValue()} miles")
index = routing.Start(0)
plan_output = "Route for vehicle 0:\n"
route_distance = 0
while not routing.IsEnd(index):
plan_output += f" {manager.IndexToNode(index) + 1} ->"
previous_index = index
index = solution.Value(routing.NextVar(index))
route_distance += routing.GetArcCostForVehicle(previous_index, index, 0)
plan_output += f" {manager.IndexToNode(index)}\n"
plan_output += f"Route distance: {route_distance} miles\n"
print(plan_output)
def main():
"""Entry point of the program."""
# Instantiate the data problem.
data = create_data_model()
# Create the routing index manager.
manager = pywrapcp.RoutingIndexManager(
len(data["distance_matrix"]), data["num_vehicles"], data["depot"]
)
# Create Routing Model.
routing = pywrapcp.RoutingModel(manager)
def distance_callback(from_index, to_index):
"""Returns the distance between the two nodes."""
# Convert from routing variable Index to distance matrix NodeIndex.
from_node = manager.IndexToNode(from_index)
to_node = manager.IndexToNode(to_index)
return data["distance_matrix"][from_node][to_node]
transit_callback_index = routing.RegisterTransitCallback(distance_callback)
# Define cost of each arc.
routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)
# Setting first solution heuristic.
# Local Optimum
# search_parameters = pywrapcp.DefaultRoutingSearchParameters()
# search_parameters.first_solution_strategy = (
# routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
# )
# Global Optimum
search_parameters = pywrapcp.DefaultRoutingSearchParameters()
search_parameters.local_search_metaheuristic = (
routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
search_parameters.time_limit.seconds = 30
search_parameters.log_search = True
# Solve the problem.
solution = routing.SolveWithParameters(search_parameters)
# Print solution on console.
if solution:
print_solution(manager, routing, solution)
if __name__ == "__main__":
start = timeit.default_timer()
main()
stop = timeit.default_timer()
print('Time: ', stop - start)