import math import random import struct import zlib from dataclasses import dataclass from typing import List, Sequence, Tuple 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 _write_png(filename: str, pixels: list[list[tuple[int, int, int]]]) -> None: height = len(pixels) width = len(pixels[0]) if height else 0 def chunk(chunk_type: bytes, data: bytes) -> bytes: return ( struct.pack(">I", len(data)) + chunk_type + data + struct.pack(">I", zlib.crc32(chunk_type + data) & 0xFFFFFFFF) ) raw = b"".join(b"\x00" + bytes([c for px in row for c in px]) for row in pixels) png = b"\x89PNG\r\n\x1a\n" ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0) png += chunk(b"IHDR", ihdr) png += chunk(b"IDAT", zlib.compress(raw, 9)) png += chunk(b"IEND", b"") with open(filename, "wb") as f: f.write(png) def _scale_points(points: Sequence[tuple[float, float]], size: int = 800, margin: int = 20): xs = [p[0] for p in points] ys = [p[1] for p in points] min_x, max_x = min(xs), max(xs) min_y, max_y = min(ys), max(ys) scale_x = (size - 2 * margin) / (max_x - min_x + 1e-9) scale_y = (size - 2 * margin) / (max_y - min_y + 1e-9) return [ ( int((x - min_x) * scale_x + margin), int((y - min_y) * scale_y + margin), ) for x, y in points ] def _draw_line(pixels: list[list[tuple[int, int, int]]], p1: tuple[int, int], p2: tuple[int, int], color: tuple[int, int, int]): x1, y1 = p1 x2, y2 = p2 dx = abs(x2 - x1) dy = -abs(y2 - y1) sx = 1 if x1 < x2 else -1 sy = 1 if y1 < y2 else -1 err = dx + dy while True: if 0 <= x1 < len(pixels[0]) and 0 <= y1 < len(pixels): pixels[y1][x1] = color if x1 == x2 and y1 == y2: break e2 = 2 * err if e2 >= dy: err += dy x1 += sx if e2 <= dx: err += dx y1 += sy def _draw_circle(pixels: list[list[tuple[int, int, int]]], center: tuple[int, int], radius: int, color: tuple[int, int, int]): cx, cy = center for y in range(cy - radius, cy + radius + 1): for x in range(cx - radius, cx + radius + 1): if 0 <= x < len(pixels[0]) and 0 <= y < len(pixels): if (x - cx) ** 2 + (y - cy) ** 2 <= radius ** 2: pixels[y][x] = color def plot_tour(cities: Sequence[City], tour: Sequence[int], save_path: str) -> None: ordered = [cities[i] for i in tour] + [cities[tour[0]]] points = _scale_points(ordered) width = height = 820 pixels = [[(255, 255, 255) for _ in range(width)] for _ in range(height)] for i in range(len(points) - 1): _draw_line(pixels, points[i], points[i + 1], (0, 120, 200)) # draw cities city_points = _scale_points(cities) for p in city_points: _draw_circle(pixels, p, 4, (200, 50, 50)) _write_png(save_path, pixels) def plot_history(best_lengths: Sequence[float], save_path: str) -> None: if not best_lengths: return width, height, margin = 820, 400, 20 pixels = [[(255, 255, 255) for _ in range(width)] for _ in range(height)] n = len(best_lengths) min_len, max_len = min(best_lengths), max(best_lengths) span = max_len - min_len if max_len != min_len else 1 def to_point(idx: int, value: float) -> tuple[int, int]: x = margin + int((width - 2 * margin) * idx / max(1, n - 1)) y = height - margin - int((height - 2 * margin) * (value - min_len) / span) return x, y prev = to_point(0, best_lengths[0]) for i, v in enumerate(best_lengths[1:], start=1): cur = to_point(i, v) _draw_line(pixels, prev, cur, (30, 30, 30)) prev = cur _write_png(save_path, pixels) @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()