UKO01_PROG_TSP_NN_KH.py

# Samplescript UKO01.PROG Nearest Neighbor K-Heuristik
import math
import pytest


def calculate_distance(coord1, coord2):
    return math.sqrt((coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2)


def nearest_neighbor_tsp(start_node, nodes):
    unvisited = set(nodes.keys())
    tour = [start_node]
    unvisited.remove(start_node)

    total_distance = 0.0
    current_node = start_node

    while unvisited:
        next_node = min(unvisited, key=lambda node: calculate_distance(nodes[current_node], nodes[node]))
        distance = calculate_distance(nodes[current_node], nodes[next_node])
        total_distance += distance
        tour.append(next_node)
        unvisited.remove(next_node)
        current_node = next_node

    # Add distance back to the start to complete the tour
    total_distance += calculate_distance(nodes[current_node], nodes[start_node])

    return tour, total_distance


def read_tsp_file(file_path):
    """
    read TSPlib input file and return a nodes-dictionary

    :param file_path:
    :return: dict containing nodes with tupels representing x, y coordinates
    """
    nodes = {}
    with open(file_path, 'r') as file:
        lines = file.readlines()

        node_section = False
        for line in lines:
            line = line.strip()
            if line.startswith("NODE_COORD_SECTION"):
                node_section = True
                continue
            if line.startswith("EOF"):
                break
            if node_section:
                parts = line.split()
                if len(parts) == 3:
                    node_id, x, y = parts
                    nodes[int(node_id)] = (float(x), float(y))

    return nodes


# *******************************************************************************************************************
# * Testcases to validate if code is working as expected                                                            *
# *******************************************************************************************************************

def test_tsp_test_1():
    """
    Test, whether test4.tsp returns a tour with length 80

    :return:
    """
    file_path = './data/test4.tsp'
    nodes = read_tsp_file(file_path)
    start_node = 1
    _, total_distance = nearest_neighbor_tsp(start_node, nodes)

    assert total_distance == 80.0, f"Expected distance to be 80, but got {total_distance}"


def test_tsp_test_2():
    """
    Test, whether test8.tsp returns a tour with length 4*5 + 4*5*sqrt(2) = 20*(1+sqrt(2))=48.28

    :return:
    """
    file_path = './data/test8.tsp'
    nodes = read_tsp_file(file_path)
    start_node = 1
    tour, total_distance = nearest_neighbor_tsp(start_node, nodes)

    # make sure, tour follows the points in counter clockwise direction onthe octagon
    assert tour == [1, 2, 3, 4, 5, 6, 7, 8], f"Expected tour to be the sequence of 1-8, but got {tour}"
    # make sure, tour distance summary equals to the 4x5 units pieces + 4x5xsqrt(2) units pieces
    assert total_distance == pytest.approx(4 * 5 + 4 * 5 * math.sqrt(2),
                                           rel=1e-4), f"Expected distance to be 48.2842, but got {total_distance}."


# *******************************************************************************************************************
# * Main Program                                                                                                    *
# *******************************************************************************************************************


def main():
    file_paths = ['./data/berlin52.tsp', './data/u1817.tsp', './data/u2319.tsp']
    start_nodes = [1, 257, 2319]
    optimal_cost = [7542, 57201, 234256]

    for id, file in enumerate(file_paths):
        nodes = read_tsp_file(file)
        tour, dist = nearest_neighbor_tsp(start_nodes[id], nodes)

        print(f"*** File {file}, with start node {id}:")
        print(f"Tour: {tour[:5]}")
        print(f"Dist: {round(dist)}")
        print(f"Opt.Gap: {round(100 * (dist / optimal_cost[id] - 1), 2)}")


if __name__ == "__main__":
    main()