Skip to content

Commit b74d7f5

Browse files
Siddhant231xyzSiddhant231xyz
committed
Modified doctest
1 parent 06485a2 commit b74d7f5

File tree

1 file changed

+104
-114
lines changed

1 file changed

+104
-114
lines changed

‎data_structures/binary_tree/lowest_common_ancestor.py

+104-114
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,16 @@
22
# https://en.wikipedia.org/wiki/Breadth-first_search
33

44
from __future__ import annotations
5-
65
from queue import Queue
76

87

98
def swap(a: int, b: int) -> tuple[int, int]:
109
"""
11-
Return a tuple (b, a) when given two integers a and b
12-
>>> swap(2,3)
10+
Return a tuple (b, a) when given two integers a and b.
11+
12+
>>> swap(2, 3)
1313
(3, 2)
14-
>>> swap(3,4)
14+
>>> swap(3, 4)
1515
(4, 3)
1616
>>> swap(67, 12)
1717
(12, 67)
@@ -24,22 +24,30 @@ def swap(a: int, b: int) -> tuple[int, int]:
2424

2525
def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
2626
"""
27-
Create a sparse table which saves each node's 2^i-th parent.
28-
29-
>>> max_node = 5
30-
>>> parent = [
31-
... [0, 0, 1, 1, 2, 2], # 2^0-th parents
32-
... [0, 0, 0, 0, 1, 1] # 2^1-th parents
33-
... ]
34-
>>> create_sparse(max_node, parent)
35-
[[0, 0, 1, 1, 2, 2], [0, 0, 0, 0, 1, 1]]
36-
>>> max_node = 3
37-
>>> parent = [
38-
... [0, 0, 1, 1], # 2^0-th parents
39-
... [0, 0, 0, 0] # 2^1-th parents
40-
... ]
41-
>>> create_sparse(max_node, parent)
42-
[[0, 0, 1, 1], [0, 0, 0, 0]]
27+
Create a sparse table that saves each node's 2^i-th parent.
28+
29+
The given `parent` table should have the direct parent of each node in row 0.
30+
The function then fills in parent[j][i] = parent[j-1][parent[j-1][i]] for each j where 2^j < max_node.
31+
32+
For example, consider a small tree where:
33+
- Node 1 is the root (its parent is 0),
34+
- Nodes 2 and 3 have parent 1.
35+
36+
We set up the parent table for only two levels (row 0 and row 1)
37+
for max_node = 3. (Note that in practice the table has many rows.)
38+
39+
>>> # Create an initial parent table with 2 rows and indices 0..3.
40+
>>> parent0 = [0, 0, 1, 1] # 0 is unused; node1's parent=0, node2 and 3's parent=1.
41+
>>> parent1 = [0, 0, 0, 0]
42+
>>> parent = [parent0, parent1]
43+
>>> # We need at least (1 << j) < max_node holds only for j = 1 here since (1 << 1)=2 < 3 and (1 << 2)=4 !< 3.
44+
>>> sparse = create_sparse(3, parent)
45+
>>> sparse[1][1], sparse[1][2], sparse[1][3]
46+
(0, 0, 0)
47+
>>> # Explanation:
48+
>>> # For node 1: parent[1][1] = parent[0][parent[0][1]] = parent[0][0] = 0.
49+
>>> # For node 2: parent[1][2] = parent[0][parent[0][2]] = parent[0][1] = 0.
50+
>>> # For node 3: parent[1][3] = parent[0][parent[0][3]] = parent[0][1] = 0.
4351
"""
4452
j = 1
4553
while (1 << j) < max_node:
@@ -49,69 +57,46 @@ def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
4957
return parent
5058

5159

52-
# returns lca of node u,v
5360
def lowest_common_ancestor(
5461
u: int, v: int, level: list[int], parent: list[list[int]]
5562
) -> int:
5663
"""
57-
Return the lowest common ancestor of nodes u and v.
58-
59-
>>> max_node = 13
60-
>>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
61-
>>> level = [-1 for _ in range(max_node + 10)]
62-
>>> graph = {
63-
... 1: [2, 3, 4],
64-
... 2: [5],
65-
... 3: [6, 7],
66-
... 4: [8],
67-
... 5: [9, 10],
68-
... 6: [11],
69-
... 7: [],
70-
... 8: [12, 13],
71-
... 9: [],
72-
... 10: [],
73-
... 11: [],
74-
... 12: [],
75-
... 13: [],
76-
... }
77-
>>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
78-
>>> parent = create_sparse(max_node, parent)
79-
>>> lowest_common_ancestor(1, 3, level, parent)
80-
1
81-
>>> lowest_common_ancestor(5, 6, level, parent)
64+
Return the lowest common ancestor (LCA) of nodes u and v in a tree.
65+
66+
The lists `level` and `parent` must be precomputed. `level[i]` is the depth of node i,
67+
and `parent` is a sparse table where parent[0][i] is the direct parent of node i.
68+
69+
>>> # Consider a simple tree:
70+
>>> # 1
71+
>>> # / \\
72+
>>> # 2 3
73+
>>> # With levels: level[1]=0, level[2]=1, level[3]=1 and parent[0]=[0,0,1,1]
74+
>>> level = [-1, 0, 1, 1] # index 0 is dummy
75+
>>> parent = [[0, 0, 1, 1]] + [[0, 0, 0, 0] for _ in range(19)]
76+
>>> lowest_common_ancestor(2, 3, level, parent)
8277
1
83-
>>> lowest_common_ancestor(7, 11, level, parent)
84-
1
85-
>>> lowest_common_ancestor(6, 7, level, parent)
86-
3
87-
>>> lowest_common_ancestor(4, 12, level, parent)
88-
4
89-
>>> lowest_common_ancestor(8, 8, level, parent)
90-
8
91-
>>> lowest_common_ancestor(9, 10, level, parent)
92-
5
93-
>>> lowest_common_ancestor(12, 13, level, parent)
94-
8
78+
>>> # LCA of a node with itself is itself.
79+
>>> lowest_common_ancestor(2, 2, level, parent)
80+
2
9581
"""
96-
# u must be deeper in the tree than v
82+
# Ensure u is at least as deep as v.
9783
if level[u] < level[v]:
9884
u, v = swap(u, v)
99-
# making depth of u same as depth of v
85+
# Bring u up to the same level as v.
10086
for i in range(18, -1, -1):
10187
if level[u] - (1 << i) >= level[v]:
10288
u = parent[i][u]
103-
# at the same depth if u==v that mean lca is found
89+
# If they are the same, we've found the LCA.
10490
if u == v:
10591
return u
106-
# moving both nodes upwards till lca in found
92+
# Move u and v up together until the LCA is found.
10793
for i in range(18, -1, -1):
10894
if parent[i][u] not in [0, parent[i][v]]:
10995
u, v = parent[i][u], parent[i][v]
110-
# returning longest common ancestor of u,v
96+
# Return the parent (direct ancestor) which is the LCA.
11197
return parent[0][u]
11298

11399

114-
# runs a breadth first search from root node of the tree
115100
def breadth_first_search(
116101
level: list[int],
117102
parent: list[list[int]],
@@ -120,54 +105,23 @@ def breadth_first_search(
120105
root: int = 1,
121106
) -> tuple[list[int], list[list[int]]]:
122107
"""
123-
Perform a breadth-first search from the root node of the tree.
124-
Sets every node's direct parent and calculates the depth of each node from the root.
125-
126-
>>> max_node = 5
127-
>>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
128-
>>> level = [-1 for _ in range(max_node + 10)]
129-
>>> graph = {
130-
... 1: [2, 3],
131-
... 2: [4],
132-
... 3: [5],
133-
... 4: [],
134-
... 5: []
135-
... }
136-
>>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
137-
>>> level[:6]
138-
[ -1, 0, 1, 1, 2, 2]
139-
>>> parent[0][1] == 0
140-
True
141-
>>> parent[0][2] == 1
142-
True
143-
>>> parent[0][3] == 1
144-
True
145-
>>> parent[0][4] == 2
146-
True
147-
>>> parent[0][5] == 3
148-
True
149-
150-
>>> # Test with disconnected graph
151-
>>> max_node = 4
152-
>>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
153-
>>> level = [-1 for _ in range(max_node + 10)]
154-
>>> graph = {
155-
... 1: [2],
156-
... 2: [],
157-
... 3: [4],
158-
... 4: []
159-
... }
160-
>>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
161-
>>> level[:5]
162-
[ -1, 0, 1, -1, -1]
163-
>>> parent[0][1] == 0
164-
True
165-
>>> parent[0][2] == 1
166-
True
167-
>>> parent[0][3] == 0
168-
True
169-
>>> parent[0][4] == 3
170-
True
108+
Run a breadth-first search (BFS) from the root node of the tree.
109+
110+
Sets every node's direct parent (in parent[0]) and calculates the depth (level)
111+
of each node from the root.
112+
113+
>>> # Consider a simple tree:
114+
>>> # 1
115+
>>> # / \\
116+
>>> # 2 3
117+
>>> graph = {1: [2, 3], 2: [], 3: []}
118+
>>> level = [-1] * 4 # index 0 is unused; nodes 1 to 3.
119+
>>> parent = [[0] * 4 for _ in range(20)]
120+
>>> new_level, new_parent = breadth_first_search(level, parent, 3, graph, root=1)
121+
>>> new_level[1:4]
122+
[0, 1, 1]
123+
>>> new_parent[0][1:4]
124+
[0, 1, 1]
171125
"""
172126
level[root] = 0
173127
q: Queue[int] = Queue(maxsize=max_node)
@@ -183,10 +137,46 @@ def breadth_first_search(
183137

184138

185139
def main() -> None:
140+
"""
141+
Run a BFS to set node depths and parents in a sample tree,
142+
then create the sparse table and compute several lowest common ancestors.
143+
144+
The sample tree used is:
145+
146+
1
147+
/ | \
148+
2 3 4
149+
/ / \\ \\
150+
5 6 7 8
151+
/ \\ | / \\
152+
9 10 11 12 13
153+
154+
The expected lowest common ancestors are:
155+
- LCA(1, 3) --> 1
156+
- LCA(5, 6) --> 1
157+
- LCA(7, 11) --> 3
158+
- LCA(6, 7) --> 3
159+
- LCA(4, 12) --> 4
160+
- LCA(8, 8) --> 8
161+
162+
To test main() without it printing to the console, we capture the output.
163+
164+
>>> import sys
165+
>>> from io import StringIO
166+
>>> backup = sys.stdout
167+
>>> sys.stdout = StringIO()
168+
>>> main()
169+
>>> output = sys.stdout.getvalue()
170+
>>> sys.stdout = backup
171+
>>> 'LCA of node 1 and 3 is: 1' in output
172+
True
173+
>>> 'LCA of node 7 and 11 is: 3' in output
174+
True
175+
"""
186176
max_node = 13
187-
# initializing with 0
177+
# initializing with 0; extra space is allocated.
188178
parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
189-
# initializing with -1 which means every node is unvisited
179+
# initializing with -1 which means every node is unvisited.
190180
level = [-1 for _ in range(max_node + 10)]
191181
graph: dict[int, list[int]] = {
192182
1: [2, 3, 4],

0 commit comments

Comments
 (0)