Skip to content

Commit a8f05fe

Browse files
Add doctests and type hints (#10974)
* Add doctests and type hints * Apply suggestions from code review * Update tarjans_scc.py * Update tarjans_scc.py --------- Co-authored-by: Tianyi Zheng <tianyizheng02@gmail.com>
1 parent 1a5d5cf commit a8f05fe

File tree

1 file changed

+29
-6
lines changed

1 file changed

+29
-6
lines changed

graphs/tarjans_scc.py

+29-6
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from collections import deque
22

33

4-
def tarjan(g):
4+
def tarjan(g: list[list[int]]) -> list[list[int]]:
55
"""
66
Tarjan's algo for finding strongly connected components in a directed graph
77
@@ -19,15 +19,30 @@ def tarjan(g):
1919
Complexity: strong_connect() is called at most once for each node and has a
2020
complexity of O(|E|) as it is DFS.
2121
Therefore this has complexity O(|V| + |E|) for a graph G = (V, E)
22+
23+
>>> tarjan([[2, 3, 4], [2, 3, 4], [0, 1, 3], [0, 1, 2], [1]])
24+
[[4, 3, 1, 2, 0]]
25+
>>> tarjan([[], [], [], []])
26+
[[0], [1], [2], [3]]
27+
>>> a = [0, 1, 2, 3, 4, 5, 4]
28+
>>> b = [1, 0, 3, 2, 5, 4, 0]
29+
>>> n = 7
30+
>>> sorted(tarjan(create_graph(n, list(zip(a, b))))) == sorted(
31+
... tarjan(create_graph(n, list(zip(a[::-1], b[::-1])))))
32+
True
33+
>>> a = [0, 1, 2, 3, 4, 5, 6]
34+
>>> b = [0, 1, 2, 3, 4, 5, 6]
35+
>>> sorted(tarjan(create_graph(n, list(zip(a, b)))))
36+
[[0], [1], [2], [3], [4], [5], [6]]
2237
"""
2338

2439
n = len(g)
25-
stack = deque()
40+
stack: deque[int] = deque()
2641
on_stack = [False for _ in range(n)]
2742
index_of = [-1 for _ in range(n)]
2843
lowlink_of = index_of[:]
2944

30-
def strong_connect(v, index, components):
45+
def strong_connect(v: int, index: int, components: list[list[int]]) -> int:
3146
index_of[v] = index # the number when this node is seen
3247
lowlink_of[v] = index # lowest rank node reachable from here
3348
index += 1
@@ -57,16 +72,24 @@ def strong_connect(v, index, components):
5772
components.append(component)
5873
return index
5974

60-
components = []
75+
components: list[list[int]] = []
6176
for v in range(n):
6277
if index_of[v] == -1:
6378
strong_connect(v, 0, components)
6479

6580
return components
6681

6782

68-
def create_graph(n, edges):
69-
g = [[] for _ in range(n)]
83+
def create_graph(n: int, edges: list[tuple[int, int]]) -> list[list[int]]:
84+
"""
85+
>>> n = 7
86+
>>> source = [0, 0, 1, 2, 3, 3, 4, 4, 6]
87+
>>> target = [1, 3, 2, 0, 1, 4, 5, 6, 5]
88+
>>> edges = list(zip(source, target))
89+
>>> create_graph(n, edges)
90+
[[1, 3], [2], [0], [1, 4], [5, 6], [], [5]]
91+
"""
92+
g: list[list[int]] = [[] for _ in range(n)]
7093
for u, v in edges:
7194
g[u].append(v)
7295
return g

0 commit comments

Comments
 (0)