mirror of
https://github.com/TheAlgorithms/Python.git
synced 2025-02-28 11:58:39 +00:00
219 lines
6.2 KiB
Python
219 lines
6.2 KiB
Python
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)
|
|
(3, 2)
|
|
>>> swap(3, 4)
|
|
(4, 3)
|
|
>>> swap(67, 12)
|
|
(12, 67)
|
|
"""
|
|
a ^= b
|
|
b ^= a
|
|
a ^= b
|
|
return a, b
|
|
|
|
|
|
def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
|
|
"""
|
|
Create a sparse table that saves each node's 2^i-th parent.
|
|
|
|
The given ``parent`` table should have the direct parent of each node in row 0.
|
|
This function fills in:
|
|
|
|
parent[j][i] = parent[j - 1][parent[j - 1][i]]
|
|
|
|
for each j where 2^j is less than 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.)
|
|
|
|
>>> parent0 = [0, 0, 1, 1] # 0 is unused; node1's parent=0, nodes 2 and 3's parent=1.
|
|
>>> parent1 = [0, 0, 0, 0]
|
|
>>> parent = [parent0, parent1]
|
|
>>> sparse = create_sparse(3, parent)
|
|
>>> (sparse[1][1], sparse[1][2], sparse[1][3])
|
|
(0, 0, 0)
|
|
"""
|
|
j = 1
|
|
while (1 << j) < max_node:
|
|
for i in range(1, max_node + 1):
|
|
parent[j][i] = parent[j - 1][parent[j - 1][i]]
|
|
j += 1
|
|
return parent
|
|
|
|
|
|
def lowest_common_ancestor(
|
|
u: int, v: int, level: list[int], parent: list[list[int]]
|
|
) -> int:
|
|
"""
|
|
Return the lowest common ancestor (LCA) of nodes u and v in a tree.
|
|
<<<<<<< HEAD
|
|
|
|
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.
|
|
|
|
=======
|
|
|
|
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.
|
|
|
|
>>>>>>> 097e9c6149e80f095be1b3dbef1c04ff94a7325a
|
|
>>> # 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(2, 2, level, parent)
|
|
2
|
|
"""
|
|
# Ensure u is at least as deep as v.
|
|
if level[u] < level[v]:
|
|
u, v = swap(u, 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]
|
|
# If they are the same, we've found the LCA.
|
|
if u == v:
|
|
return u
|
|
# 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]
|
|
return parent[0][u]
|
|
|
|
|
|
def breadth_first_search(
|
|
level: list[int],
|
|
parent: list[list[int]],
|
|
max_node: int,
|
|
graph: dict[int, list[int]],
|
|
root: int = 1,
|
|
) -> tuple[list[int], list[list[int]]]:
|
|
"""
|
|
Run a breadth-first search (BFS) from the root node of the tree.
|
|
|
|
This sets each node's direct parent (stored in parent[0]) and calculates the
|
|
depth (level) of each node from the root.
|
|
|
|
>>> # 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)
|
|
q.put(root)
|
|
while q.qsize() != 0:
|
|
u = q.get()
|
|
for v in graph[u]:
|
|
if level[v] == -1:
|
|
level[v] = level[u] + 1
|
|
q.put(v)
|
|
parent[0][v] = u
|
|
return level, parent
|
|
|
|
|
|
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:
|
|
<<<<<<< HEAD
|
|
|
|
1
|
|
/ | \
|
|
2 3 4
|
|
/ / \\ \\
|
|
5 6 7 8
|
|
/ \\ | / \\
|
|
9 10 11 12 13
|
|
|
|
=======
|
|
|
|
1
|
|
/ | \
|
|
2 3 4
|
|
/ / \\ \\
|
|
5 6 7 8
|
|
/ \\ | / \\
|
|
9 10 11 12 13
|
|
|
|
>>>>>>> 097e9c6149e80f095be1b3dbef1c04ff94a7325a
|
|
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
|
|
parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
|
|
level = [-1 for _ in range(max_node + 10)]
|
|
graph: dict[int, list[int]] = {
|
|
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)
|
|
print("LCA of node 1 and 3 is: ", lowest_common_ancestor(1, 3, level, parent))
|
|
print("LCA of node 5 and 6 is: ", lowest_common_ancestor(5, 6, level, parent))
|
|
print("LCA of node 7 and 11 is: ", lowest_common_ancestor(7, 11, level, parent))
|
|
print("LCA of node 6 and 7 is: ", lowest_common_ancestor(6, 7, level, parent))
|
|
print("LCA of node 4 and 12 is: ", lowest_common_ancestor(4, 12, level, parent))
|
|
print("LCA of node 8 and 8 is: ", lowest_common_ancestor(8, 8, level, parent))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|