51 lines
2.3 KiB
Python

import math, random, time, numpy as np
def distance(city1, city2):
return math.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2) + 1e-10
def total_distance(cities):
return sum([distance(cities[i - 1], cities[i]) for i in range(len(cities))])
class AntColony:
def __init__(self, cities, n_ants, alpha=1, beta=2, evaporation=0.5, intensification=2, max_time=0.1):
self.cities = cities
self.n = len(cities)
self.n_ants = n_ants
self.alpha = alpha
self.beta = beta
self.evaporation = evaporation
self.intensification = intensification
self.max_time = max_time
self.pheromones = [[1 / self.n for _ in range(self.n)] for __ in range(self.n)]
def choose_next_city(self, ant):
unvisited_cities = [i for i in range(self.n) if i not in ant]
probabilities = [self.pheromones[ant[-1]][i] ** self.alpha * ((1 / distance(self.cities[ant[-1]], self.cities[i])) ** self.beta) for i in unvisited_cities]
total = sum(probabilities)
if total == 0:
probabilities = [1 / len(unvisited_cities) for _ in unvisited_cities]
else:
probabilities = [p / total for p in probabilities]
return np.random.choice(unvisited_cities, p=probabilities)
def update_pheromones(self, ant):
pheromones_delta = self.intensification / total_distance([self.cities[i] for i in ant])
for i in range(len(ant) - 1):
self.pheromones[ant[i]][ant[i+1]] += pheromones_delta
def run(self):
best_ant = []
best_distance = float('inf')
start_time = time.time()
while time.time() - start_time < self.max_time:
ants = [[random.randint(0, self.n - 1)] for _ in range(self.n_ants)]
for ant in ants:
for _ in range(self.n - 1):
ant.append(self.choose_next_city(ant))
ant_distance = total_distance([self.cities[i] for i in ant])
if ant_distance < best_distance:
best_distance = ant_distance
best_ant = ant.copy()
self.update_pheromones(ant)
self.pheromones = [[(1 - self.evaporation) * p for p in row] for row in self.pheromones]
return [self.cities[i] for i in best_ant]