import numpy as np
import matplotlib.pyplot as plt

# 定义问题
n = 10  # 城市数量
distance = np.array([[0, 2, 9, 10, 10, 15, 20, 25, 30, 35],
                     [2, 0, 6,  7,  8,  12, 17, 22, 27, 32],
                      [9, 6, 0,  1,  2,  6,  11, 16, 21, 26],
                     [10, 7, 1,  0,  1,  5,  10, 15, 20, 25],
                     [10, 8, 2,  1,  0,  4,  9,  14, 19, 24],
                     [15, 12, 6, 5, 4, 0, 5, 10, 15, 20],
                     [20, 17, 11, 10, 9, 5, 0, 5, 10, 15],
                     [25, 22, 16, 15, 14, 10, 5, 0, 5, 10],
                     [30, 27, 21, 20, 19, 15, 10, 5, 0, 5],
                     [35, 32, 26, 25, 24, 20, 15, 10, 5, 0]])  # 城市间距离矩阵

# 初始化参数
m = 50  # 蚂蚁数量
alpha = 1  # 信息素重要程度因子
beta = 2  # 启发式因子
rho = 0.5  # 信息素挥发因子
Q = 100  # 常数

# 初始化信息素矩阵
pheromone = np.ones((n, n))
# 迭代优化
max_iter = 100  # 最大迭代次数
best_path = []  # 最佳路径
best_distance = np.inf  # 最佳距离

for iter in range(max_iter):
    # 构建解空间
    path = np.zeros((m, n), dtype=int)
    path[:, 0] = np.random.permutation(n)  # 随机初始化起点

    # 计算每只蚂蚁的路径距离和信息素更新量
    delta_pheromone = np.zeros((n, n))
    for i in range(m):
        # 逐城市选择下一个城市
        for j in range(1, n):
            # 计算当前城市到其他城市的转移概率
            prob = (pheromone[path[i, j-1], :] ** alpha) * ((1 / distance[path[i, j-1], :]) ** beta)
            prob[path[i, :j]] = 0  # 已访问的城市不可再访问
            prob /= prob.sum()
            # 根据转移概率选择下一个城市
            path[i, j] = np.random.choice(n, p=prob)
            # 更新信息素
        delta_pheromone += Q / distance[path[i, :n-1], path[i, 1:]]

    # 更新全局最优解
    for i in range(m):
        if distance[path[i, -1], path[i, 0]] + distance[path[i, 0], path[i, 1:-1]].sum() < best_distance:
            best_distance = distance[path[i, -1], path[i, 0]] + distance[path[i, 0], path[i, 1:-1]].sum()
            best_path = path[i]

# 更新信息素
pheromone = (1 - rho) * pheromone + delta_pheromone

# 输出结果
print("最佳路径:", best_path)
print("最佳距离:", best_distance)
plt.plot(best_path)
plt.show()