Modified doctest

This commit is contained in:
Siddhant Jain 2025-01-13 16:59:00 -05:00
parent b74d7f5392
commit d6fff7504f

View File

@ -1,6 +1,3 @@
# https://en.wikipedia.org/wiki/Lowest_common_ancestor
# https://en.wikipedia.org/wiki/Breadth-first_search
from __future__ import annotations
from queue import Queue
@ -26,8 +23,12 @@ 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.
The function then fills in parent[j][i] = parent[j-1][parent[j-1][i]] for each j where 2^j < max_node.
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),
@ -36,18 +37,12 @@ def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
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.
>>> 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]
>>> # 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]
>>> (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:
@ -63,19 +58,19 @@ def lowest_common_ancestor(
"""
Return the lowest common ancestor (LCA) of nodes u and v in a tree.
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.
>>> # 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]
>>> # 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
>>> # LCA of a node with itself is itself.
>>> lowest_common_ancestor(2, 2, level, parent)
2
"""
@ -93,7 +88,6 @@ def lowest_common_ancestor(
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 the parent (direct ancestor) which is the LCA.
return parent[0][u]
@ -107,8 +101,8 @@ def breadth_first_search(
"""
Run a breadth-first search (BFS) from the root node of the tree.
Sets every node's direct parent (in parent[0]) and calculates the depth (level)
of each node from the root.
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
@ -138,8 +132,8 @@ 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.
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:
@ -174,9 +168,7 @@ def main() -> None:
True
"""
max_node = 13
# 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.
level = [-1 for _ in range(max_node + 10)]
graph: dict[int, list[int]] = {
1: [2, 3, 4],