mirror of
https://github.com/TheAlgorithms/Python.git
synced 2024-11-27 23:11:09 +00:00
143 lines
4.5 KiB
Python
143 lines
4.5 KiB
Python
"""
|
|
Finding the shortest path in 0-1-graph in O(E + V) which is faster than dijkstra.
|
|
0-1-graph is the weighted graph with the weights equal to 0 or 1.
|
|
Link: https://codeforces.com/blog/entry/22276
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from collections.abc import Iterator
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class Edge:
|
|
"""Weighted directed graph edge."""
|
|
|
|
destination_vertex: int
|
|
weight: int
|
|
|
|
|
|
class AdjacencyList:
|
|
"""Graph adjacency list."""
|
|
|
|
def __init__(self, size: int):
|
|
self._graph: list[list[Edge]] = [[] for _ in range(size)]
|
|
self._size = size
|
|
|
|
def __getitem__(self, vertex: int) -> Iterator[Edge]:
|
|
"""Get all the vertices adjacent to the given one."""
|
|
return iter(self._graph[vertex])
|
|
|
|
@property
|
|
def size(self):
|
|
return self._size
|
|
|
|
def add_edge(self, from_vertex: int, to_vertex: int, weight: int):
|
|
"""
|
|
>>> g = AdjacencyList(2)
|
|
>>> g.add_edge(0, 1, 0)
|
|
>>> g.add_edge(1, 0, 1)
|
|
>>> list(g[0])
|
|
[Edge(destination_vertex=1, weight=0)]
|
|
>>> list(g[1])
|
|
[Edge(destination_vertex=0, weight=1)]
|
|
>>> g.add_edge(0, 1, 2)
|
|
Traceback (most recent call last):
|
|
...
|
|
ValueError: Edge weight must be either 0 or 1.
|
|
>>> g.add_edge(0, 2, 1)
|
|
Traceback (most recent call last):
|
|
...
|
|
ValueError: Vertex indexes must be in [0; size).
|
|
"""
|
|
if weight not in (0, 1):
|
|
raise ValueError("Edge weight must be either 0 or 1.")
|
|
|
|
if to_vertex < 0 or to_vertex >= self.size:
|
|
raise ValueError("Vertex indexes must be in [0; size).")
|
|
|
|
self._graph[from_vertex].append(Edge(to_vertex, weight))
|
|
|
|
def get_shortest_path(self, start_vertex: int, finish_vertex: int) -> int | None:
|
|
"""
|
|
Return the shortest distance from start_vertex to finish_vertex in 0-1-graph.
|
|
1 1 1
|
|
0--------->3 6--------7>------->8
|
|
| ^ ^ ^ |1
|
|
| | | |0 v
|
|
0| |0 1| 9-------->10
|
|
| | | ^ 1
|
|
v | | |0
|
|
1--------->2<-------4------->5
|
|
0 1 1
|
|
>>> g = AdjacencyList(11)
|
|
>>> g.add_edge(0, 1, 0)
|
|
>>> g.add_edge(0, 3, 1)
|
|
>>> g.add_edge(1, 2, 0)
|
|
>>> g.add_edge(2, 3, 0)
|
|
>>> g.add_edge(4, 2, 1)
|
|
>>> g.add_edge(4, 5, 1)
|
|
>>> g.add_edge(4, 6, 1)
|
|
>>> g.add_edge(5, 9, 0)
|
|
>>> g.add_edge(6, 7, 1)
|
|
>>> g.add_edge(7, 8, 1)
|
|
>>> g.add_edge(8, 10, 1)
|
|
>>> g.add_edge(9, 7, 0)
|
|
>>> g.add_edge(9, 10, 1)
|
|
>>> g.add_edge(1, 2, 2)
|
|
Traceback (most recent call last):
|
|
...
|
|
ValueError: Edge weight must be either 0 or 1.
|
|
>>> g.get_shortest_path(0, 3)
|
|
0
|
|
>>> g.get_shortest_path(0, 4)
|
|
Traceback (most recent call last):
|
|
...
|
|
ValueError: No path from start_vertex to finish_vertex.
|
|
>>> g.get_shortest_path(4, 10)
|
|
2
|
|
>>> g.get_shortest_path(4, 8)
|
|
2
|
|
>>> g.get_shortest_path(0, 1)
|
|
0
|
|
>>> g.get_shortest_path(1, 0)
|
|
Traceback (most recent call last):
|
|
...
|
|
ValueError: No path from start_vertex to finish_vertex.
|
|
"""
|
|
queue = deque([start_vertex])
|
|
distances: list[int | None] = [None] * self.size
|
|
distances[start_vertex] = 0
|
|
|
|
while queue:
|
|
current_vertex = queue.popleft()
|
|
current_distance = distances[current_vertex]
|
|
if current_distance is None:
|
|
continue
|
|
|
|
for edge in self[current_vertex]:
|
|
new_distance = current_distance + edge.weight
|
|
dest_vertex_distance = distances[edge.destination_vertex]
|
|
if (
|
|
isinstance(dest_vertex_distance, int)
|
|
and new_distance >= dest_vertex_distance
|
|
):
|
|
continue
|
|
distances[edge.destination_vertex] = new_distance
|
|
if edge.weight == 0:
|
|
queue.appendleft(edge.destination_vertex)
|
|
else:
|
|
queue.append(edge.destination_vertex)
|
|
|
|
if distances[finish_vertex] is None:
|
|
raise ValueError("No path from start_vertex to finish_vertex.")
|
|
|
|
return distances[finish_vertex]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
|
|
doctest.testmod()
|