Pure Python implementations of competitive programming algorithms.
ckp.graph_theory.tree
contains implementations of algorithms on graphs.
Unless noted otherwise, all tree-related functions don’t recurse, so the stack will not overflow even for large trees!
Tree-related algorithms revolve around the following data structure:
class TreeData:
root: int
parents: list[int]
children: list[list[int]]
depths: list[int]
sizes: list[int]|None
def __len__(self) -> int: pass
def __repr__(self) -> int: pass
Each node of a tree is represented by an integer between 0 and len(tree)-1
.
tree.root
: The root node of a tree. Typically 0
, but could be any other node.tree.parents
: Parent of each node. The root node’s parent is -1
.tree.children
: Children of each node.tree.depths
: Depth of each node. The root node’s depth is 0
.tree.sizes
: Size of subtree of each node.
tree_sizes(tree)
automatically sets tree.sizes
if it does not exist.
def tree_from_neighbors(neighbors: list[list[int]], root: int = 0) -> TreeData
Create a TreeData
, given a bi-directional adjacency list.
def tree_from_parents(parents: list[int], root: int = 0) -> TreeData
Create a TreeData
, given a list of parents.
def tree_from_edges(edges: list[tuple[int, int]], root: int = 0) -> TreeData
Create a TreeData
, given a list of edges. This is probably the most common way to create a tree.
def tree_with_root(tree: TreeData, new_root: int) -> TreeData
Re-root the given tree.
def tree_sizes(tree: TreeData) -> list[int]
Return tree.sizes
, size of subtree for each node. If it’s missing, then compute the list and cache it to tree
.
def tree_centroids(tree: TreeData) -> list[int]
Returns the centroid (or centroids) of the tree.
A centroid is a a vertex whose largest subtree has at most len(tree) // 2
vertices.
A tree either has one or two centroids, so this function always return a list of one or two vertices (as long as the tree is non-empty).
def tree_lca_init(tree: TreeData) -> TreeLCAData
Pre-compute ancestor information for tree
.
def tree_lca_pth_ancestor(lca: TreeLCAData, v: int, p: int) -> int
Return the p
-th ancestor of v
.
def tree_lca_query(lca: TreeLCAData, v: int, w: int) -> int
Return the lowest common ancestor of v
and w
using the pre-computed data.
class TreeHLDData:
path_index: list[tuple[int, int]]
paths: list[list[int]]
def tree_hld_init(tree: TreeData) -> TreeHLDData
def tree_hld_decompose_descendant(hld: TreeHLDData, node: int, ancestor: int)
def euler_tour(tree: TreeData) -> EulerTourData
Return an Euler tour of tree
. tour.visits
lists the nodes in DFS order and tour.begin[v]
gives the index of the first visit to v
.
def euler_tour_sorted(tree: TreeData) -> EulerTourData
Similar to euler_tour
but children are visited in an order suitable for heavy-light decomposition so that nodes in the same heavy path appear consecutively in the tour.
WARNING: The current implementation of tree isomorphism related functions do recurse, and very inefficient (while ensuring linear runtime).
def rooted_tree_isomorphism_signature(tree: TreeData) -> int
def rooted_tree_is_isomorphic_to(a: TreeData, b: TreeData) -> bool
Compute an isomorphism signature for a rooted tree and test whether two rooted trees are isomorphic.
def tree_isomorphism_signature(tree: TreeData) -> int
def tree_is_isomorphic_to(a: TreeData, b: TreeData) -> bool
Functions for un-rooted trees. Centroids are chosen automatically and their signatures are compared.