161 lines
4.6 KiB
Python
161 lines
4.6 KiB
Python
import math
|
|
import random
|
|
from dataclasses import dataclass
|
|
from typing import List, Sequence, Tuple
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
City = Tuple[float, float]
|
|
Tour = List[int]
|
|
|
|
|
|
def euclidean_distance(c1: City, c2: City) -> float:
|
|
return math.hypot(c1[0] - c2[0], c1[1] - c2[1])
|
|
|
|
|
|
def build_distance_matrix(cities: Sequence[City]) -> list[list[float]]:
|
|
size = len(cities)
|
|
matrix = [[0.0 for _ in range(size)] for _ in range(size)]
|
|
for i in range(size):
|
|
for j in range(i + 1, size):
|
|
dist = euclidean_distance(cities[i], cities[j])
|
|
matrix[i][j] = matrix[j][i] = dist
|
|
return matrix
|
|
|
|
|
|
def plot_tour(cities: Sequence[City], tour: Sequence[int], save_path: str) -> None:
|
|
x = [cities[i][0] for i in tour]
|
|
y = [cities[i][1] for i in tour]
|
|
|
|
fig, ax = plt.subplots(figsize=(7, 7))
|
|
ax.plot(x + [x[0]], y + [y[0]], "k-", linewidth=1)
|
|
ax.plot(x, y, "ro", markersize=4)
|
|
|
|
ax.axis("equal")
|
|
fig.tight_layout()
|
|
fig.savefig(save_path, dpi=220)
|
|
plt.close(fig)
|
|
|
|
|
|
def plot_history(best_lengths: Sequence[float], save_path: str) -> None:
|
|
if not best_lengths:
|
|
return
|
|
|
|
iterations = list(range(len(best_lengths)))
|
|
|
|
fig, ax = plt.subplots(figsize=(10, 6))
|
|
ax.plot(iterations, best_lengths, linewidth=2, color="blue")
|
|
|
|
ax.set_xlabel("Итерация", fontsize=12)
|
|
ax.set_ylabel("Длина лучшего тура", fontsize=12)
|
|
ax.grid(True, alpha=0.3)
|
|
|
|
fig.savefig(save_path, dpi=150, bbox_inches="tight")
|
|
plt.close(fig)
|
|
|
|
|
|
@dataclass
|
|
class ACOConfig:
|
|
cities: Sequence[City]
|
|
n_ants: int
|
|
n_iterations: int
|
|
alpha: float = 1.0
|
|
beta: float = 5.0
|
|
rho: float = 0.5
|
|
q: float = 1.0
|
|
seed: int | None = None
|
|
|
|
|
|
@dataclass
|
|
class ACOResult:
|
|
best_tour: Tour
|
|
best_length: float
|
|
history: List[float]
|
|
|
|
|
|
class AntColonyOptimizer:
|
|
def __init__(self, config: ACOConfig):
|
|
self.config = config
|
|
if config.seed is not None:
|
|
random.seed(config.seed)
|
|
|
|
self.cities = config.cities
|
|
self.dist_matrix = build_distance_matrix(config.cities)
|
|
n = len(config.cities)
|
|
self.pheromone = [[1.0 if i != j else 0.0 for j in range(n)] for i in range(n)]
|
|
|
|
def _choose_next_city(self, current: int, unvisited: set[int]) -> int:
|
|
candidates = list(unvisited)
|
|
weights = []
|
|
for nxt in candidates:
|
|
tau = self.pheromone[current][nxt] ** self.config.alpha
|
|
eta = (1.0 / (self.dist_matrix[current][nxt] + 1e-12)) ** self.config.beta
|
|
weights.append(tau * eta)
|
|
|
|
total = sum(weights)
|
|
probs = [w / total for w in weights]
|
|
return random.choices(candidates, weights=probs, k=1)[0]
|
|
|
|
def _build_tour(self, start: int) -> Tour:
|
|
n = len(self.cities)
|
|
tour = [start]
|
|
unvisited = set(range(n))
|
|
unvisited.remove(start)
|
|
|
|
current = start
|
|
while unvisited:
|
|
nxt = self._choose_next_city(current, unvisited)
|
|
tour.append(nxt)
|
|
unvisited.remove(nxt)
|
|
current = nxt
|
|
|
|
return tour
|
|
|
|
def _tour_length(self, tour: Sequence[int]) -> float:
|
|
return sum(
|
|
self.dist_matrix[tour[i]][tour[(i + 1) % len(tour)]]
|
|
for i in range(len(tour))
|
|
)
|
|
|
|
def run(self) -> ACOResult:
|
|
best_tour: Tour = []
|
|
best_length = float("inf")
|
|
best_history: list[float] = []
|
|
|
|
for _ in range(self.config.n_iterations):
|
|
tours: list[Tour] = []
|
|
lengths: list[float] = []
|
|
|
|
for _ in range(self.config.n_ants):
|
|
start_city = random.randrange(len(self.cities))
|
|
tour = self._build_tour(start_city)
|
|
length = self._tour_length(tour)
|
|
tours.append(tour)
|
|
lengths.append(length)
|
|
|
|
if length < best_length:
|
|
best_length = length
|
|
best_tour = tour
|
|
|
|
for i in range(len(self.pheromone)):
|
|
for j in range(len(self.pheromone)):
|
|
self.pheromone[i][j] *= 1 - self.config.rho
|
|
|
|
for tour, length in zip(tours, lengths):
|
|
deposit = self.config.q / length
|
|
for i in range(len(tour)):
|
|
a, b = tour[i], tour[(i + 1) % len(tour)]
|
|
self.pheromone[a][b] += deposit
|
|
self.pheromone[b][a] += deposit
|
|
|
|
best_history.append(best_length)
|
|
|
|
return ACOResult(
|
|
best_tour=best_tour, best_length=best_length, history=best_history
|
|
)
|
|
|
|
|
|
def run_aco(config: ACOConfig) -> ACOResult:
|
|
optimizer = AntColonyOptimizer(config)
|
|
return optimizer.run()
|