Python中使用蚁群算法解决TSP问题

import random
import math

class Ant:
    def __init__(self, num_cities):
        self.num_cities = num_cities
        self.path = []
        self.visited = [False] * num_cities
        self.distance = 0.0

    def get_distance(self, distances):
        for i in range(self.num_cities):
            self.distance += distances[self.path[i]][self.path[(i+1)%self.num_cities]]

    def clear(self):
        self.path = []
        self.visited = [False] * self.num_cities
        self.distance = 0.0

class AntColony:
    def __init__(self, num_ants, num_cities, distances, alpha, beta, rho, q):
        self.num_ants = num_ants
        self.num_cities = num_cities
        self.distances = distances
        self.alpha = alpha
        self.beta = beta
        self.rho = rho
        self.q = q
        self.pheromones = [[1.0/self.num_cities] * self.num_cities for i in range(self.num_cities)]
        self.ants = [Ant(self.num_cities) for i in range(self.num_ants)]
        self.best_ant = Ant(self.num_cities)

    def update_pheromones(self):
        for i in range(self.num_cities):
            for j in range(self.num_cities):
                self.pheromones[i][j] *= (1.0 - self.rho)
        for ant in self.ants:
            for i in range(self.num_cities):
                j = (i+1)%self.num_cities
                self.pheromones[ant.path[i]][ant.path[j]] += self.q/ant.distance

    def select_next_city(self, ant):
        i = ant.path[-1]
        probabilities = []
        total_prob = 0.0
        for j in range(self.num_cities):
            if not ant.visited[j]:
                tau = self.pheromones[i][j]**self.alpha
                eta = (1.0/self.distances[i][j])**self.beta
                probabilities.append(tau * eta)
                total_prob += tau * eta
            else:
                probabilities.append(0.0)
        if total_prob > 0.0:
            probabilities = [p/total_prob for p in probabilities]
            next_city = random.choices(range(self.num_cities), weights=probabilities)[0]
        else:
            next_city = random.choice([j for j in range(self.num_cities) if not ant.visited[j]])
        ant.path.append(next_city)
        ant.visited[next_city] = True

    def run(self, max_iterations):
        for it in range(max_iterations):
            for ant in self.ants:
                ant.clear()
                start_city = random.randint(0, self.num_cities-1)
                ant.path.append(start_city)
                ant.visited[start_city] = True
                for i in range(self.num_cities-1):
                    self.select_next_city(ant)
                ant.get_distance(self.distances)
                if ant.distance < self.best_ant.distance or self.best_ant.distance == 0.0:
                    self.best_ant = ant
            self.update_pheromones()

        return self.best_ant


if __name__ == '__main__':
    num_ants = 10
    num_cities = 10
    distances = [[0.0]*num_cities for i in range(num_cities)]
    for i in range(num_cities):
        for j in range(num_cities):
            if i == j:
                distances[i][j] = 0.0
            else:
                distances[i][j] = distances[j][i] = math.sqrt((i-j)**2 + (i+j)**2)

    alpha = 1.0
    beta = 3.0
    rho = 0.1
    q = 10.0
    ant_colony = AntColony(num_ants, num_cities, distances, alpha, beta, rho, q)
    best_ant = ant_colony.run(1000)
    print('Best path:', best_ant.path)
    print('Best distance:', best_ant.distance)