johnson_algorithm.py

"""
Реализация алгоритма Джонсона для поиска кратчайших путей
между всеми парами вершин в графе.
"""

from typing import List, Tuple, Optional, Dict
from collections import defaultdict
import heapq
import graphviz


class Graph:
    """Представление графа для алгоритма Джонсона"""

    def __init__(self, num_vertices: int):
        self.num_vertices = num_vertices
        self.edges: List[Tuple[int, int, float]] = []  # (from, to, weight)
        self.adjacency_matrix: Optional[List[List[float]]] = None

    def add_edge(self, from_vertex: int, to_vertex: int, weight: float):
        """Добавление ребра в граф"""
        if 0 <= from_vertex < self.num_vertices and 0 <= to_vertex < self.num_vertices:
            self.edges.append((from_vertex, to_vertex, weight))
        else:
            raise ValueError(
                f"Вершины должны быть в диапазоне [0, {self.num_vertices - 1}]"
            )

    def from_adjacency_matrix(self, matrix: List[List[float]]):
        """Создание графа из матрицы смежности"""
        if len(matrix) != self.num_vertices:
            raise ValueError(
                f"Размер матрицы должен быть {self.num_vertices}x{self.num_vertices}"
            )

        self.adjacency_matrix = matrix
        self.edges = []

        for i in range(self.num_vertices):
            for j in range(self.num_vertices):
                # Добавляем ребро, если оно существует (не inf)
                # Вес 0 - это валидный вес ребра (не петля), поэтому включаем его
                # Петли (i == j) с весом 0 также учитываем для полноты
                if matrix[i][j] != float("inf"):
                    self.edges.append((i, j, matrix[i][j]))


def bellman_ford(graph: Graph, source: int) -> Tuple[bool, List[float]]:
    """
    Алгоритм Беллмана-Форда для поиска кратчайших путей от источника.
    Возвращает (найден_отрицательный_цикл, расстояния).
    """
    distances = [float("inf")] * graph.num_vertices
    distances[source] = 0

    # Релаксация ребер (V-1) раз
    for _ in range(graph.num_vertices - 1):
        for from_v, to_v, weight in graph.edges:
            if distances[from_v] != float("inf"):
                if distances[from_v] + weight < distances[to_v]:
                    distances[to_v] = distances[from_v] + weight

    # Проверка на отрицательные циклы
    for from_v, to_v, weight in graph.edges:
        if distances[from_v] != float("inf"):
            if distances[from_v] + weight < distances[to_v]:
                return True, distances  # Найден отрицательный цикл

    return False, distances


def dijkstra(graph: Graph, source: int, h: List[float]) -> List[float]:
    """
    Алгоритм Дейкстры для поиска кратчайших путей от источника.
    Использует перевешенные веса h.
    """
    # Перевешивание ребер
    reweighted_edges: Dict[int, List[Tuple[int, float]]] = defaultdict(list)
    for from_v, to_v, weight in graph.edges:
        new_weight = weight + h[from_v] - h[to_v]
        reweighted_edges[from_v].append((to_v, new_weight))

    # Инициализация
    distances = [float("inf")] * graph.num_vertices
    distances[source] = 0
    visited = [False] * graph.num_vertices
    heap = [(0, source)]

    while heap:
        current_dist, current_vertex = heapq.heappop(heap)

        if visited[current_vertex]:
            continue

        visited[current_vertex] = True

        # Обновление расстояний до соседей
        if current_vertex in reweighted_edges:
            for neighbor, weight in reweighted_edges[current_vertex]:
                if not visited[neighbor]:
                    new_dist = current_dist + weight
                    if new_dist < distances[neighbor]:
                        distances[neighbor] = new_dist
                        heapq.heappush(heap, (new_dist, neighbor))

    # Восстановление исходных расстояний
    for i in range(graph.num_vertices):
        if distances[i] != float("inf"):
            distances[i] = distances[i] - h[source] + h[i]

    return distances


def johnson_algorithm(
    graph: Graph,
) -> Tuple[bool, Optional[List[List[float]]], Optional[str]]:
    """
    Алгоритм Джонсона для поиска кратчайших путей между всеми парами вершин.

    Возвращает:
        (успех, матрица_расстояний, сообщение_об_ошибке)
    """
    # Шаг 1: Создание нового графа с фиктивной вершиной
    extended_graph = Graph(graph.num_vertices + 1)

    # Копируем все ребра исходного графа
    for from_v, to_v, weight in graph.edges:
        extended_graph.add_edge(from_v, to_v, weight)

    # Добавляем фиктивную вершину (последняя) с нулевыми ребрами ко всем остальным
    fake_vertex = graph.num_vertices
    for i in range(graph.num_vertices):
        extended_graph.add_edge(fake_vertex, i, 0)

    # Шаг 2: Применение Беллмана-Форда для перевешивания
    has_negative_cycle, h = bellman_ford(extended_graph, fake_vertex)

    if has_negative_cycle:
        return False, None, "Обнаружен отрицательный цикл в графе!"

    # Удаляем фиктивную вершину из h
    h = h[:-1]

    # Шаг 3: Применение Дейкстры для каждой вершины
    result_matrix = []
    for source in range(graph.num_vertices):
        distances = dijkstra(graph, source, h)
        result_matrix.append(distances)

    return True, result_matrix, None


def parse_adjacency_matrix(text: str) -> Tuple[int, List[List[float]]]:
    """
    Парсинг матрицы смежности из текста.
    Формат: числа разделены пробелами, строки разделены переносами строк.
    'inf' или пустая строка означает отсутствие ребра.
    """
    lines = [line.strip() for line in text.strip().split("\n") if line.strip()]

    if not lines:
        raise ValueError("Матрица пуста")

    matrix = []
    num_vertices = len(lines)

    for i, line in enumerate(lines):
        values = line.split()
        if len(values) != num_vertices:
            raise ValueError(
                f"Строка {i + 1}: ожидается {num_vertices} значений, получено {len(values)}"
            )

        row = []
        for j, val in enumerate(values):
            val = val.strip().lower()
            if val == "" or val == "inf" or val == "infinity" or val == "∞":
                row.append(float("inf"))
            else:
                try:
                    row.append(float(val))
                except ValueError:
                    raise ValueError(
                        f"Строка {i + 1}, столбец {j + 1}: неверное значение '{val}'"
                    )

        matrix.append(row)

    return num_vertices, matrix


def format_result_matrix(
    matrix: List[List[float]], vertex_names: Optional[List[str]] = None
) -> str:
    """Форматирование матрицы результатов для вывода"""
    if not matrix:
        return "Нет результатов"

    n = len(matrix)
    if vertex_names is None:
        vertex_names = [str(i) for i in range(n)]

    # Определение ширины столбцов
    col_widths = [len(name) for name in vertex_names]
    max_name_width = max(len(name) for name in vertex_names)

    for i in range(n):
        for j in range(n):
            if matrix[i][j] == float("inf"):
                width = 3  # "inf"
            else:
                width = len(f"{matrix[i][j]:.2f}")
            col_widths[j] = max(col_widths[j], width)

    # Заголовок
    result = "Матрица кратчайших расстояний:\n\n"
    result += " " * (max_name_width + 2)
    for j, name in enumerate(vertex_names):
        result += f"{name:>{max(col_widths[j], len(name))}} "
    result += (
        "\n"
        + "-" * (sum(col_widths) + len(vertex_names) * 2 + max_name_width + 2)
        + "\n"
    )

    # Данные
    for i, name in enumerate(vertex_names):
        result += f"{name:>{max_name_width}} | "
        for j in range(n):
            if matrix[i][j] == float("inf"):
                result += f"{'inf':>{col_widths[j]}} "
            else:
                # Форматируем число с правильным выравниванием
                formatted_num = f"{matrix[i][j]:.2f}"
                result += f"{formatted_num:>{col_widths[j]}} "
        result += "\n"

    return result


def find_shortest_paths_info(matrix: List[List[float]]) -> str:
    """Генерация дополнительной информации о кратчайших путях"""
    if not matrix:
        return ""

    n = len(matrix)
    info = "\n" + "=" * 50 + "\n"
    info += "Дополнительная информация:\n"
    info += "=" * 50 + "\n\n"

    # Нахождение минимального и максимального расстояний
    min_dist = float("inf")
    max_dist = -float("inf")
    min_pair = None
    max_pair = None

    for i in range(n):
        for j in range(n):
            if i != j and matrix[i][j] != float("inf"):
                if matrix[i][j] < min_dist:
                    min_dist = matrix[i][j]
                    min_pair = (i, j)
                if matrix[i][j] > max_dist:
                    max_dist = matrix[i][j]
                    max_pair = (i, j)

    if min_pair is not None:
        info += f"Кратчайший путь: {min_pair[0]} → {min_pair[1]} (расстояние: {min_dist:.2f})\n"
    if max_pair is not None:
        info += f"Длиннейший путь: {max_pair[0]} → {max_pair[1]} (расстояние: {max_dist:.2f})\n"

    # Подсчет недостижимых вершин
    unreachable_count = 0
    for i in range(n):
        for j in range(n):
            if i != j and matrix[i][j] == float("inf"):
                unreachable_count += 1

    if unreachable_count > 0:
        info += f"\nНедостижимых пар вершин: {unreachable_count}\n"
    else:
        info += "\nВсе вершины достижимы друг от друга.\n"

    return info


def visualize_graph(
    graph: Graph,
    output_file: Optional[str] = None,
    format: str = "png",
    highlight_negative_edges: bool = True,
) -> graphviz.Digraph:
    """
    Визуализация графа с помощью Graphviz.

    Args:
        graph: Граф для визуализации
        output_file: Путь к файлу для сохранения (если None, не сохраняется)
        format: Формат изображения (png, svg, pdf и т.д.)
        highlight_negative_edges: Подсветка ребер с отрицательными весами

    Returns:
        Объект Digraph для дальнейшей настройки
    """
    # Создание направленного графа
    dot = graphviz.Digraph(comment="Graph Visualization")
    dot.attr(rankdir="LR")  # Слева направо
    dot.attr("node", shape="circle", style="filled", fillcolor="lightblue")
    dot.attr("edge", fontsize="10")

    # Добавление вершин
    for i in range(graph.num_vertices):
        dot.node(str(i), str(i))

    # Добавление ребер
    for from_v, to_v, weight in graph.edges:
        # Определение цвета и стиля ребра
        edge_color = "red" if weight < 0 and highlight_negative_edges else "black"
        edge_style = "dashed" if weight < 0 and highlight_negative_edges else "solid"

        # Форматирование веса
        if weight == int(weight):
            weight_label = str(int(weight))
        else:
            weight_label = f"{weight:.2f}"

        dot.edge(
            str(from_v),
            str(to_v),
            label=weight_label,
            color=edge_color,
            style=edge_style,
        )

    # Сохранение файла, если указан путь
    if output_file:
        try:
            dot.render(output_file, format=format, cleanup=True)
        except Exception as e:
            raise RuntimeError(f"Не удалось сохранить визуализацию: {str(e)}")

    return dot