Python实现蚁群算法(样例）

import numpy as np

class AntColony:
    def __init__(self, distances, n_ants, n_best, n_iterations, decay, alpha=1, beta=1):
        """
        Args:
            distances (2D numpy.array): 2D array of distances between points
            n_ants (int): Number of ants running per iteration
            n_best (int): Number of best ants who deposit pheromone
            n_iteration (int): Number of iterations
            decay (float): Rate it which pheromone decays. The pheromone value is multiplied by decay, so 0.95 will lead to decay, 0.5 to much faster decay.
            alpha (int or float): exponenet on pheromone, higher alpha gives pheromone more weight. Default=1
            beta (int or float): exponent on distance, higher beta give distance more weight. Default=1
        """
        self.distances  = distances
        self.pheromone = np.ones(self.distances.shape) / len(distances)
        self.all_inds = range(len(distances))
        self.n_ants = n_ants
        self.n_best = n_best
        self.n_iterations = n_iterations
        self.decay = decay
        self.alpha = alpha
        self.beta = beta
        
    def run(self):
        shortest_path = None
        shortest_path_length = np.inf
        for i in range(self.n_iterations):
            all_paths = self.gen_all_paths()
            self.spread_pheronome(all_paths, self.n_best, shortest_path=shortest_path, shortest_path_length=shortest_path_length)
            shortest_path, shortest_path_length = self.pick_best_path(all_paths)
            self.pheromone * self.decay
            
        return shortest_path, shortest_path_length
            
    def spread_pheronome(self, all_paths, n_best, shortest_path, shortest_path_length):
        sorted_paths = sorted(all_paths, key=lambda x: x[1])
        for path, path_length in sorted_paths[:n_best]:
            for move in path:
                self.pheromone[move] += 1.0 / self.distances[move]
    
    def pick_best_path(self, all_paths):
        # sort all paths by length
        all_paths = sorted(all_paths, key=lambda x: x[1])
        best_path = all_paths[0]
        return best_path
    
    def gen_path_dist(self, path):
        total_dist = 0
        for ele in path:
            total_dist += self.distances[ele]
        return total_dist
    
    def gen_all_paths(self):
        all_paths = []
        for i in range(self.n_ants):
            path = self.gen_path(0)
            all_paths.append((path, self.gen_path_dist(path)))
        return all_paths
    
    def gen_path(self, start):
        path = []
        visited = set()
        visited.add(start)
        prev = start
        for i in range(len(self.distances) - 1):
            move = self.pick_move(self.pheromone[prev], self.distances[prev], visited)
            path.append((prev, move))
            prev = move
            visited.add(move)
        path.append((prev, start)) # going back to where we started    
        return path
    
    def pick_move(self, pheromone, dist, visited):
        pheromone = np.copy(pheromone)
        pheromone[list(visited)] = 0
        
        row = pheromone ** self.alpha * (( 1.0 / dist) ** self.beta)
        norm_row = row / row.sum()
        move = np_choice(self.all_inds, 1, p=norm_row)[0]
        return move

# Helper Function
def np_choice(a, size, replace=True, p=None):
    """numpy's random.choice doesn't allow for an option to not replace, 
       writing a simple function to do that"""
    idx = np.random.choice(range(len(a)), size=size, replace=replace, p=p)
    return np.array(a)[idx]

# Example Usage
if __name__ == '__main__':
    # Example distance matrix (symmetric)
    distances = np.array([[0, 2, 2, 5],
                          [2, 0, 1, 2],
                          [2, 1, 0, 1],
                          [5, 2, 1, 0]])

    # Parameters
    n_ants = 3
    n_best = 2
    n_iterations = 10
    decay = 0.95

    # Create ant colony instance
    ant_colony = AntColony(distances, n_ants, n_best, n_iterations, decay)
    
    # Running the optimization
    shortest_path, shortest_path_length = ant_colony.run()
    print("Shortest Path:", shortest_path)
    print("Shortest Path Length:", shortest_path_length)