diff --git a/pyrival/algebra/mod_template.py b/pyrival/algebra/mod_template.py new file mode 100644 index 0000000..3412d06 --- /dev/null +++ b/pyrival/algebra/mod_template.py @@ -0,0 +1,54 @@ +""" +A template for usefull stuff involving prime modulo. It contains: +1. Calculation of factorial, inverse factorial and modular inverse + for all integers < maxN in O(maxN) time. +2. Calculate n choose k in O(1) time using precalculated fac and inv_fac. +3. Multiply matrices mod MOD. +""" +MOD = 10 ** 9 + 7 # needs to be prime! +maxN = 10 ** 6 # needs to be <= MOD + +modmul = lambda a,b,c=0:(a*b + c) % MOD + +""" Precalculate factorial, inverse factorial and modular inverse """ + +def mod_precalc(n): + """ Calculates fac, inv_fac and (modular) inv for i < n in O(n) time """ + assert n <= MOD + + fac = [1] * n + for i in range(2, n): + fac[i] = modmul(fac[i - 1], i) + + inv_fac = [pow(fac[-1], MOD - 2, MOD)] * n + for i in reversed(range(1, n)): + inv_fac[i - 1] = modmul(inv_fac[i], i) + + inv = [modmul(inv_fac[i], fac[i - 1]) for i in range(n)] + + return fac, inv_fac, inv + +fac, inv_fac, inv = mod_precalc(maxN) + + +""" Useful functions involving modulo """ + +def choose(n, k): + """ Calculate n choose k in O(1) time """ + if k < 0 or k > n: + return 0 + return modmul(modmul(fac[n], inv_fac[k]), inv_fac[n - k]) + +def matrix_modmul(A, B): + """ Multiplies matrices A and B mod MOD""" + assert len(A[0]) == len(B) + C = [] + for Ai in A: + tmp = [0] * len(B[0]) + for k in range(len(B)): + Aik = Ai[k] + Bk = B[k] + for j in range(len(Bk)): + tmp[j] = modmul(Aik, Bk[j], tmp[j]) + C.append(tmp) + return C diff --git a/pyrival/graphs/centriod_shallowest_decomposition.py b/pyrival/graphs/centriod_shallowest_decomposition.py new file mode 100644 index 0000000..6e48499 --- /dev/null +++ b/pyrival/graphs/centriod_shallowest_decomposition.py @@ -0,0 +1,114 @@ +""" + This file contains the implementations of centroid decomposition + and the shallowest decomposition of a tree, based on the blog + + https://codeforces.com/blog/entry/125018 + + The two implementations are meant to be interchangable. + In particular this allows easy comparision between the + methods. + + + # EXAMPLE OF HOW TO USE (https://codeforces.com/contest/321/problem/C) + + n = int(input()) + graph = [[] for _ in range(n)] + for _ in range(n - 1): + u,v = [int(x) - 1 for x in input().split()] + graph[u].append(v) + graph[v].append(u) + + decomposition_tree, root = centroid_decomposition_tree(graph) + #decomposition_tree, root = shallowest_decomposition_tree(graph) + + labels = [0] * n + bfs = [root] + for node in bfs: + for child in decomposition_tree[node]: + labels[child] = labels[node] + 1 + bfs.append(child) + + alphabet = [chr(ord('A') + i) for i in range(26)] + print(' '.join(alphabet[label] for label in labels)) +""" + + +""" + Standard implementation of centroid decomposition using rerooting. + Takes O(n log n) time. +""" + +def centroid_decomposition_tree(graph): + n = len(graph) + graph = [c[:] for c in graph] # copy + bfs = [0] + for node in bfs: + bfs += graph[node] + for nei in graph[node]: + graph[nei].remove(node) + size = [0] * n + for node in reversed(bfs): + size[node] = 1 + sum(size[child] for child in graph[node]) + decomposition_tree = [[] for _ in range(n)] + def centroid_reroot(u): + N = size[u] + while True: + for v in graph[u]: + if size[v] > N // 2: + size[u] = N - size[v] + graph[u].remove(v) + graph[v].append(u) + u = v + break + else: # u is the centroid + decomposition_tree[u] = [centroid_reroot(v) for v in graph[u]] + return u + decomposition_root = centroid_reroot(0) + return decomposition_tree, decomposition_root + +""" + Implementation of the shallowest tree decomposition based on https://codeforces.com/blog/entry/125018 + Takes O(n) time. + This is similar to centroid decomposition except that + it sometimes is ''shallower'' than the centroid decomposition. +""" + +def ctz(x): # Count trailing zeroes + return (x & -x).bit_length() - 1 + +def shallowest_decomposition_tree(graph, root=0): + n = len(graph) + forbidden = [0] * n + decomposition_tree = [[] for _ in range(n)] + stacks = [[] for _ in range(n.bit_length())] + def extract_chain(labels, u): + while labels: + label = labels.bit_length() - 1 + labels ^= 2**label + v = stacks[label].pop() + decomposition_tree[u].append(v) + u = v + dfs = [root] + while dfs: + u = dfs.pop() + if u >= 0: + forbidden[u] = -1 + dfs.append(~u) + for v in graph[u]: + if not forbidden[v]: + dfs.append(v) + else: + u = ~u + forbidden_once = forbidden_twice = 0 + for v in graph[u]: + forbidden_twice |= forbidden_once & (forbidden[v] + 1) + forbidden_once |= forbidden[v] + 1 + forbidden[u] = forbidden_once | (2**forbidden_twice.bit_length() - 1) + label_u = ctz(forbidden[u] + 1) + stacks[label_u].append(u) + for v in graph[u]: + extract_chain((forbidden[v] + 1) & (2**label_u - 1), u) + max_label = (forbidden[root] + 1).bit_length() - 1 + decomposition_root = stacks[max_label].pop() + extract_chain((forbidden[root] + 1) & (2**max_label - 1), decomposition_root) + return decomposition_tree, decomposition_root