Skip to content

Commit 5be77f3

Browse files
authored
Add 0-1-bfs. (#3285)
* Add 0-1-bfs. * fixup! Add 0-1-bfs. * fixup! Add 0-1-bfs. * Check edge weights. * Check edge vertecies.
1 parent 89e8dbf commit 5be77f3

File tree

1 file changed

+138
-0
lines changed

1 file changed

+138
-0
lines changed

graphs/bfs_zero_one_shortest_path.py

+138
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
from collections import deque
2+
from dataclasses import dataclass
3+
from typing import Iterator, List
4+
5+
"""
6+
Finding the shortest path in 0-1-graph in O(E + V) which is faster than dijkstra.
7+
0-1-graph is the weighted graph with the weights equal to 0 or 1.
8+
Link: https://codeforces.com/blog/entry/22276
9+
"""
10+
11+
12+
@dataclass
13+
class Edge:
14+
"""Weighted directed graph edge."""
15+
16+
destination_vertex: int
17+
weight: int
18+
19+
20+
class AdjacencyList:
21+
"""Graph adjacency list."""
22+
23+
def __init__(self, size: int):
24+
self._graph: List[List[Edge]] = [[] for _ in range(size)]
25+
self._size = size
26+
27+
def __getitem__(self, vertex: int) -> Iterator[Edge]:
28+
"""Get all the vertices adjacent to the given one."""
29+
return iter(self._graph[vertex])
30+
31+
@property
32+
def size(self):
33+
return self._size
34+
35+
def add_edge(self, from_vertex: int, to_vertex: int, weight: int):
36+
"""
37+
>>> g = AdjacencyList(2)
38+
>>> g.add_edge(0, 1, 0)
39+
>>> g.add_edge(1, 0, 1)
40+
>>> list(g[0])
41+
[Edge(destination_vertex=1, weight=0)]
42+
>>> list(g[1])
43+
[Edge(destination_vertex=0, weight=1)]
44+
>>> g.add_edge(0, 1, 2)
45+
Traceback (most recent call last):
46+
...
47+
ValueError: Edge weight must be either 0 or 1.
48+
>>> g.add_edge(0, 2, 1)
49+
Traceback (most recent call last):
50+
...
51+
ValueError: Vertex indexes must be in [0; size).
52+
"""
53+
if weight not in (0, 1):
54+
raise ValueError("Edge weight must be either 0 or 1.")
55+
56+
if to_vertex < 0 or to_vertex >= self.size:
57+
raise ValueError("Vertex indexes must be in [0; size).")
58+
59+
self._graph[from_vertex].append(Edge(to_vertex, weight))
60+
61+
def get_shortest_path(self, start_vertex: int, finish_vertex: int) -> int:
62+
"""
63+
Return the shortest distance from start_vertex to finish_vertex in 0-1-graph.
64+
1 1 1
65+
0--------->3 6--------7>------->8
66+
| ^ ^ ^ |1
67+
| | | |0 v
68+
0| |0 1| 9-------->10
69+
| | | ^ 1
70+
v | | |0
71+
1--------->2<-------4------->5
72+
0 1 1
73+
>>> g = AdjacencyList(11)
74+
>>> g.add_edge(0, 1, 0)
75+
>>> g.add_edge(0, 3, 1)
76+
>>> g.add_edge(1, 2, 0)
77+
>>> g.add_edge(2, 3, 0)
78+
>>> g.add_edge(4, 2, 1)
79+
>>> g.add_edge(4, 5, 1)
80+
>>> g.add_edge(4, 6, 1)
81+
>>> g.add_edge(5, 9, 0)
82+
>>> g.add_edge(6, 7, 1)
83+
>>> g.add_edge(7, 8, 1)
84+
>>> g.add_edge(8, 10, 1)
85+
>>> g.add_edge(9, 7, 0)
86+
>>> g.add_edge(9, 10, 1)
87+
>>> g.add_edge(1, 2, 2)
88+
Traceback (most recent call last):
89+
...
90+
ValueError: Edge weight must be either 0 or 1.
91+
>>> g.get_shortest_path(0, 3)
92+
0
93+
>>> g.get_shortest_path(0, 4)
94+
Traceback (most recent call last):
95+
...
96+
ValueError: No path from start_vertex to finish_vertex.
97+
>>> g.get_shortest_path(4, 10)
98+
2
99+
>>> g.get_shortest_path(4, 8)
100+
2
101+
>>> g.get_shortest_path(0, 1)
102+
0
103+
>>> g.get_shortest_path(1, 0)
104+
Traceback (most recent call last):
105+
...
106+
ValueError: No path from start_vertex to finish_vertex.
107+
"""
108+
queue = deque([start_vertex])
109+
distances = [None for i in range(self.size)]
110+
distances[start_vertex] = 0
111+
112+
while queue:
113+
current_vertex = queue.popleft()
114+
current_distance = distances[current_vertex]
115+
116+
for edge in self[current_vertex]:
117+
new_distance = current_distance + edge.weight
118+
if (
119+
distances[edge.destination_vertex] is not None
120+
and new_distance >= distances[edge.destination_vertex]
121+
):
122+
continue
123+
distances[edge.destination_vertex] = new_distance
124+
if edge.weight == 0:
125+
queue.appendleft(edge.destination_vertex)
126+
else:
127+
queue.append(edge.destination_vertex)
128+
129+
if distances[finish_vertex] is None:
130+
raise ValueError("No path from start_vertex to finish_vertex.")
131+
132+
return distances[finish_vertex]
133+
134+
135+
if __name__ == "__main__":
136+
import doctest
137+
138+
doctest.testmod()

0 commit comments

Comments
 (0)