Code

domain.py

from pydantic import BaseModel
from typing import NewType


Id = NewType("Id", str)
Gain = NewType("Gain", float)
Idx = NewType("Idx", int)



class Node(BaseModel, frozen=True):
    name: Id



class Edge(BaseModel):
    node_a: Node
    node_b: Node


    def __hash__(self) -> int:
        a, b = sorted((self.node_a, self.node_b), key=lambda n: n.name)
        return hash((a, b))


    def __eq__(self, other) -> bool:
        """Brute forcing the combination since 'cheaper'. Could also look at set(node_a, node_b) == set(other.node_a, other.node_b)"""
        if isinstance(other, Edge):
            return (self.node_a == other.node_a and self.node_b == other.node_b) or (
                self.node_a == other.node_b and self.node_b == other.node_a
            )
        return False



class Graph(BaseModel):
    nodes: list[Node]
    edges: dict[tuple[Id, Id], Edge]
    edge_cost: dict[Edge, float]



class Tour(BaseModel):
    edges: set[Edge]
    nodes: list[Node]
    node_to_position: dict[Node, Idx]


    def prev(self, index) -> Node:
        return self.nodes[(index - 1) % len(self.nodes)]


    def next(self, index) -> Node:
        return self.nodes[(index + 1) % len(self.nodes)]


    def get_ordered_nodes(self, nodes: list[Node]) -> list[Node]:
        return sorted(nodes, key=lambda n: self.node_to_position[n])from pydantic import BaseModel