This code is a solution for CodeChef's Tree MEX problem:
Minimum excludant (or MEX for short) of a collection of integers is the smallest non-negative integer not present in the set.
You have a tree \$ T \$ with \$n\$ vertices. Consider an ordering \$ P=(v_1,\ldots,v_n) \$ of vertices of \$T\$. We construct a sequence \$A(P)=(a_1,\ldots,a_n)\$ using the following process:
- Set all \$a_i=−1\$.
- Process vertices in order \$v_1,\ldots,v_n\$. For the current vertex \$v_i \$ set \$a_i=\operatorname{MEX}(a_{u_1},\ldots,a_{u_k})\$, where \$u_1,\ldots,u_k\$ is the set of neighbours of \$v_i\$.
For instance, let \$n=3\$ and \$T\$ be the tree with edges \$(1,2)\$ and \$(2,3)\$. Then, for the ordering \$P=(1,2,3)\$ we obtain the sequence \$A(P)=(0,1,0)\$, while for the ordering \$P=(2,3,1)\$ we obtain \$A(P)=(1,0,1)\$.
Consider all \$n!\$ orders \$P\$. How many different sequences \$A(P)\$ can we obtain? Print the answer modulo \$10^9+7\$.
I have created a graph from list of tuples and then calculated mex values for each permutation.
Ex: 3 # 3 vertices
1 2 # (1,2) edge
2 3 # (2,3) edge
permutations(1,2,3) ==> we get 6 combinations
For 6 combinations we will get 6 values. I need to print distinct count of that 6 values.
Code:
# https://www.codechef.com/problems/TREEMX
from collections import defaultdict
from itertools import permutations
class Graph:
""" Graph is data structure(directed). """
def __init__(self, connections):
""" Initializing graph with set as default value. """
self._graph = defaultdict(set)
self.add_connections(connections)
def add_connections(self, connections):
""" Add connections(list of tupules) to graph. """
for node1, node2 in connections:
self.add(node1,node2)
def add(self, node1, node2):
""" Add node1 and node2 to graph which is initialized with set by default. """
self._graph[node1].add(node2)
self._graph[node2].add(node1)
def get_graph(self):
return dict(self._graph)
def mex(arr_set):
mex = 0
while mex in arr_set:
mex+=1
return mex
def process(graph, order):
a_p = [-1] * len(order)
for el in order:
a_p[el-1] = mex([a_p[u-1] for u in graph[el]])
return a_p
t = int(input())
for _ in range(t):
v = int(input())
e = []
for i in range(v-1):
e.append(tuple(map(int, input().split())))
g = Graph(e)
all_vertices = {s for i in e for s in i}
result = []
for p in permutations(all_vertices, v):
out = process(g.get_graph(), p)
result.append(out) if out not in result else None
print(len(result) % ((10**9)+7))
Constraints:
- \$1≤T≤10\$
- \$1≤n≤10^5\$
- \$1≤u_i,v_i≤n\$
- \$u_i≠v_i\$
How can I optimize the code, to get rid of the "time limit exceeded" error?
