OR Tools이용해서 TSP 문제 풀어보기

import numpy as np from ortools.constraint_solver import routing_enums_pb2 from ortools.constraint_solver import pywrapcp import matplotlib.pyplot as plt def create_distance_matrix(locations): num_locations = len(locations) distance_matrix = [] for i in range(num_locations): ilist = [] for j in range(num_locations): ilist.append(0) distance_matrix.append(ilist) for i in range(num_locations): for j in range(num_locations): distance_matrix[i][j] = int(np.sqrt((locations[i][0] - locations[j][0]) ** 2 + (locations[i][1] - locations[j][1]) ** 2)) return distance_matrix def create_data_model(locations): data = {} data["distance_matrix"] = create_distance_matrix(locations) data["num_vehicles"] = 1 data["depot"] = 0 return data def print_solution(manager, routing, solution): index = routing.Start(0) plan_output = "" route = [manager.IndexToNode(index)] route_distance = 0 while not routing.IsEnd(index): previous_index = index index = solution.Value(routing.NextVar(index)) route.append(manager.IndexToNode(index)) route_distance += routing.GetArcCostForVehicle(previous_index, index, 0) plan_output += f" {manager.IndexToNode(previous_index)} ->" plan_output += f" {manager.IndexToNode(index)}\n" print(plan_output) print(f"거리: {route_distance}") return route, route_distance def plot_solution(locations, route, route_distance): x = locations[:, 0] y = locations[:, 1] plt.figure(figsize=(10, 10)) plt.scatter(x, y, color='red') for i in range(len(route) - 1): start_node = route[i] end_node = route[i + 1] plt.plot([locations[start_node][0], locations[end_node][0]], [locations[start_node][1], locations[end_node][1]], color='blue', linestyle='-', linewidth=2) plt.xlabel('X') plt.ylabel('Y') plt.title('TSP Google OR Tools: '+str(route_distance)) plt.legend() plt.grid(True) plt.show() def main(): locations = np.array([[0.0, 0.0], [0.0, 900], [100, 500], [200, 200], [400, 100], [400, 800], [700, 200], [800, 500], [900, 0.0], [900, 900], [0.0, 100], [0.0, 700], [400, 0.0], [400, 100], [400, 800], [400, 900], [700, 0.0], [700, 900], [900, 200], [900, 700]]) data = create_data_model(locations) manager = pywrapcp.RoutingIndexManager( len(data["distance_matrix"]), data["num_vehicles"], data["depot"] ) # 라우팅 모델 생성 routing = pywrapcp.RoutingModel(manager) def distance_callback(from_index, to_index): 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) # 비용 정의 routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index) # 첫 번째 해법 탐색 전략 설정 search_parameters = pywrapcp.DefaultRoutingSearchParameters() search_parameters.first_solution_strategy = ( routing_enums_pb2.FirstSolutionStrategy.SAVINGS ) # 문제 해결 solution = routing.SolveWithParameters(search_parameters) # 해법 출력 및 시각화 if solution: route, route_distance = print_solution(manager, routing, solution) plot_solution(locations, route, route_distance) if __name__ == "__main__": main()