mirror of
https://github.com/TheAlgorithms/Python.git
synced 2025-03-03 13:28:40 +00:00
Modified doctest
This commit is contained in:
Siddhant Jain
parent
06485a2f01
commit
b74d7f5392
@ -2,16 +2,16 @@
|
||||
# https://en.wikipedia.org/wiki/Breadth-first_search
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from queue import Queue
|
||||
|
||||
|
||||
def swap(a: int, b: int) -> tuple[int, int]:
|
||||
"""
|
||||
Return a tuple (b, a) when given two integers a and b
|
||||
>>> swap(2,3)
|
||||
Return a tuple (b, a) when given two integers a and b.
|
||||
|
||||
>>> swap(2, 3)
|
||||
(3, 2)
|
||||
>>> swap(3,4)
|
||||
>>> swap(3, 4)
|
||||
(4, 3)
|
||||
>>> swap(67, 12)
|
||||
(12, 67)
|
||||
@ -24,22 +24,30 @@ def swap(a: int, b: int) -> tuple[int, int]:
|
||||
|
||||
def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
|
||||
"""
|
||||
Create a sparse table which saves each node's 2^i-th parent.
|
||||
Create a sparse table that saves each node's 2^i-th parent.
|
||||
|
||||
>>> max_node = 5
|
||||
>>> parent = [
|
||||
... [0, 0, 1, 1, 2, 2], # 2^0-th parents
|
||||
... [0, 0, 0, 0, 1, 1] # 2^1-th parents
|
||||
... ]
|
||||
>>> create_sparse(max_node, parent)
|
||||
[[0, 0, 1, 1, 2, 2], [0, 0, 0, 0, 1, 1]]
|
||||
>>> max_node = 3
|
||||
>>> parent = [
|
||||
... [0, 0, 1, 1], # 2^0-th parents
|
||||
... [0, 0, 0, 0] # 2^1-th parents
|
||||
... ]
|
||||
>>> create_sparse(max_node, parent)
|
||||
[[0, 0, 1, 1], [0, 0, 0, 0]]
|
||||
The given `parent` table should have the direct parent of each node in row 0.
|
||||
The function then fills in parent[j][i] = parent[j-1][parent[j-1][i]] for each j where 2^j < max_node.
|
||||
|
||||
For example, consider a small tree where:
|
||||
- Node 1 is the root (its parent is 0),
|
||||
- Nodes 2 and 3 have parent 1.
|
||||
|
||||
We set up the parent table for only two levels (row 0 and row 1)
|
||||
for max_node = 3. (Note that in practice the table has many rows.)
|
||||
|
||||
>>> # Create an initial parent table with 2 rows and indices 0..3.
|
||||
>>> parent0 = [0, 0, 1, 1] # 0 is unused; node1's parent=0, node2 and 3's parent=1.
|
||||
>>> parent1 = [0, 0, 0, 0]
|
||||
>>> parent = [parent0, parent1]
|
||||
>>> # We need at least (1 << j) < max_node holds only for j = 1 here since (1 << 1)=2 < 3 and (1 << 2)=4 !< 3.
|
||||
>>> sparse = create_sparse(3, parent)
|
||||
>>> sparse[1][1], sparse[1][2], sparse[1][3]
|
||||
(0, 0, 0)
|
||||
>>> # Explanation:
|
||||
>>> # For node 1: parent[1][1] = parent[0][parent[0][1]] = parent[0][0] = 0.
|
||||
>>> # For node 2: parent[1][2] = parent[0][parent[0][2]] = parent[0][1] = 0.
|
||||
>>> # For node 3: parent[1][3] = parent[0][parent[0][3]] = parent[0][1] = 0.
|
||||
"""
|
||||
j = 1
|
||||
while (1 << j) < max_node:
|
||||
@ -49,69 +57,46 @@ def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
|
||||
return parent
|
||||
|
||||
|
||||
# returns lca of node u,v
|
||||
def lowest_common_ancestor(
|
||||
u: int, v: int, level: list[int], parent: list[list[int]]
|
||||
) -> int:
|
||||
"""
|
||||
Return the lowest common ancestor of nodes u and v.
|
||||
Return the lowest common ancestor (LCA) of nodes u and v in a tree.
|
||||
|
||||
>>> max_node = 13
|
||||
>>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
|
||||
>>> level = [-1 for _ in range(max_node + 10)]
|
||||
>>> graph = {
|
||||
... 1: [2, 3, 4],
|
||||
... 2: [5],
|
||||
... 3: [6, 7],
|
||||
... 4: [8],
|
||||
... 5: [9, 10],
|
||||
... 6: [11],
|
||||
... 7: [],
|
||||
... 8: [12, 13],
|
||||
... 9: [],
|
||||
... 10: [],
|
||||
... 11: [],
|
||||
... 12: [],
|
||||
... 13: [],
|
||||
... }
|
||||
>>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
|
||||
>>> parent = create_sparse(max_node, parent)
|
||||
>>> lowest_common_ancestor(1, 3, level, parent)
|
||||
The lists `level` and `parent` must be precomputed. `level[i]` is the depth of node i,
|
||||
and `parent` is a sparse table where parent[0][i] is the direct parent of node i.
|
||||
|
||||
>>> # Consider a simple tree:
|
||||
>>> # 1
|
||||
>>> # / \\
|
||||
>>> # 2 3
|
||||
>>> # With levels: level[1]=0, level[2]=1, level[3]=1 and parent[0]=[0,0,1,1]
|
||||
>>> level = [-1, 0, 1, 1] # index 0 is dummy
|
||||
>>> parent = [[0, 0, 1, 1]] + [[0, 0, 0, 0] for _ in range(19)]
|
||||
>>> lowest_common_ancestor(2, 3, level, parent)
|
||||
1
|
||||
>>> lowest_common_ancestor(5, 6, level, parent)
|
||||
1
|
||||
>>> lowest_common_ancestor(7, 11, level, parent)
|
||||
1
|
||||
>>> lowest_common_ancestor(6, 7, level, parent)
|
||||
3
|
||||
>>> lowest_common_ancestor(4, 12, level, parent)
|
||||
4
|
||||
>>> lowest_common_ancestor(8, 8, level, parent)
|
||||
8
|
||||
>>> lowest_common_ancestor(9, 10, level, parent)
|
||||
5
|
||||
>>> lowest_common_ancestor(12, 13, level, parent)
|
||||
8
|
||||
>>> # LCA of a node with itself is itself.
|
||||
>>> lowest_common_ancestor(2, 2, level, parent)
|
||||
2
|
||||
"""
|
||||
# u must be deeper in the tree than v
|
||||
# Ensure u is at least as deep as v.
|
||||
if level[u] < level[v]:
|
||||
u, v = swap(u, v)
|
||||
# making depth of u same as depth of v
|
||||
# Bring u up to the same level as v.
|
||||
for i in range(18, -1, -1):
|
||||
if level[u] - (1 << i) >= level[v]:
|
||||
u = parent[i][u]
|
||||
# at the same depth if u==v that mean lca is found
|
||||
# If they are the same, we've found the LCA.
|
||||
if u == v:
|
||||
return u
|
||||
# moving both nodes upwards till lca in found
|
||||
# Move u and v up together until the LCA is found.
|
||||
for i in range(18, -1, -1):
|
||||
if parent[i][u] not in [0, parent[i][v]]:
|
||||
u, v = parent[i][u], parent[i][v]
|
||||
# returning longest common ancestor of u,v
|
||||
# Return the parent (direct ancestor) which is the LCA.
|
||||
return parent[0][u]
|
||||
|
||||
|
||||
# runs a breadth first search from root node of the tree
|
||||
def breadth_first_search(
|
||||
level: list[int],
|
||||
parent: list[list[int]],
|
||||
@ -120,54 +105,23 @@ def breadth_first_search(
|
||||
root: int = 1,
|
||||
) -> tuple[list[int], list[list[int]]]:
|
||||
"""
|
||||
Perform a breadth-first search from the root node of the tree.
|
||||
Sets every node's direct parent and calculates the depth of each node from the root.
|
||||
Run a breadth-first search (BFS) from the root node of the tree.
|
||||
|
||||
>>> max_node = 5
|
||||
>>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
|
||||
>>> level = [-1 for _ in range(max_node + 10)]
|
||||
>>> graph = {
|
||||
... 1: [2, 3],
|
||||
... 2: [4],
|
||||
... 3: [5],
|
||||
... 4: [],
|
||||
... 5: []
|
||||
... }
|
||||
>>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
|
||||
>>> level[:6]
|
||||
[ -1, 0, 1, 1, 2, 2]
|
||||
>>> parent[0][1] == 0
|
||||
True
|
||||
>>> parent[0][2] == 1
|
||||
True
|
||||
>>> parent[0][3] == 1
|
||||
True
|
||||
>>> parent[0][4] == 2
|
||||
True
|
||||
>>> parent[0][5] == 3
|
||||
True
|
||||
Sets every node's direct parent (in parent[0]) and calculates the depth (level)
|
||||
of each node from the root.
|
||||
|
||||
>>> # Test with disconnected graph
|
||||
>>> max_node = 4
|
||||
>>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
|
||||
>>> level = [-1 for _ in range(max_node + 10)]
|
||||
>>> graph = {
|
||||
... 1: [2],
|
||||
... 2: [],
|
||||
... 3: [4],
|
||||
... 4: []
|
||||
... }
|
||||
>>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
|
||||
>>> level[:5]
|
||||
[ -1, 0, 1, -1, -1]
|
||||
>>> parent[0][1] == 0
|
||||
True
|
||||
>>> parent[0][2] == 1
|
||||
True
|
||||
>>> parent[0][3] == 0
|
||||
True
|
||||
>>> parent[0][4] == 3
|
||||
True
|
||||
>>> # Consider a simple tree:
|
||||
>>> # 1
|
||||
>>> # / \\
|
||||
>>> # 2 3
|
||||
>>> graph = {1: [2, 3], 2: [], 3: []}
|
||||
>>> level = [-1] * 4 # index 0 is unused; nodes 1 to 3.
|
||||
>>> parent = [[0] * 4 for _ in range(20)]
|
||||
>>> new_level, new_parent = breadth_first_search(level, parent, 3, graph, root=1)
|
||||
>>> new_level[1:4]
|
||||
[0, 1, 1]
|
||||
>>> new_parent[0][1:4]
|
||||
[0, 1, 1]
|
||||
"""
|
||||
level[root] = 0
|
||||
q: Queue[int] = Queue(maxsize=max_node)
|
||||
@ -183,10 +137,46 @@ def breadth_first_search(
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Run a BFS to set node depths and parents in a sample tree,
|
||||
then create the sparse table and compute several lowest common ancestors.
|
||||
|
||||
The sample tree used is:
|
||||
|
||||
1
|
||||
/ | \
|
||||
2 3 4
|
||||
/ / \\ \\
|
||||
5 6 7 8
|
||||
/ \\ | / \\
|
||||
9 10 11 12 13
|
||||
|
||||
The expected lowest common ancestors are:
|
||||
- LCA(1, 3) --> 1
|
||||
- LCA(5, 6) --> 1
|
||||
- LCA(7, 11) --> 3
|
||||
- LCA(6, 7) --> 3
|
||||
- LCA(4, 12) --> 4
|
||||
- LCA(8, 8) --> 8
|
||||
|
||||
To test main() without it printing to the console, we capture the output.
|
||||
|
||||
>>> import sys
|
||||
>>> from io import StringIO
|
||||
>>> backup = sys.stdout
|
||||
>>> sys.stdout = StringIO()
|
||||
>>> main()
|
||||
>>> output = sys.stdout.getvalue()
|
||||
>>> sys.stdout = backup
|
||||
>>> 'LCA of node 1 and 3 is: 1' in output
|
||||
True
|
||||
>>> 'LCA of node 7 and 11 is: 3' in output
|
||||
True
|
||||
"""
|
||||
max_node = 13
|
||||
# initializing with 0
|
||||
# initializing with 0; extra space is allocated.
|
||||
parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
|
||||
# initializing with -1 which means every node is unvisited
|
||||
# initializing with -1 which means every node is unvisited.
|
||||
level = [-1 for _ in range(max_node + 10)]
|
||||
graph: dict[int, list[int]] = {
|
||||
1: [2, 3, 4],
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user