Assignment5.py

import numpy as np

# Distance matrix
dist = np.array([[0,2,2,5],[2,0,3,4],[2,3,0,1],[5,4,1,0]])
n = len(dist)

# Parameters
ants, alpha, beta, rho, iterations = 10, 1, 2, 0.5, 50
pher = np.ones((n,n))

def prob(i, visited):
    p = (pher[i]**alpha) * ((1/(dist[i]+1e-10))**beta)
    p[list(visited)] = 0
    return p / p.sum()

best_path, best_len = None, 1e9
for _ in range(iterations):
    all_paths = []
    for _ in range(ants):
        path = [np.random.randint(n)]
        while len(path) < n:
            i = path[-1]
            next_city = np.random.choice(range(n), p=prob(i, path))
            path.append(next_city)
        path.append(path[0])
        length = sum(dist[path[i], path[i+1]] for i in range(n))
        all_paths.append((path, length))
        if length < best_len: best_path, best_len = path, length
    pher *= (1 - rho)
    for path, length in all_paths:
        for i in range(n): pher[path[i], path[i+1]] += 1/length

print("Best path:", best_path)
print("Shortest distance:", best_len)