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.


# 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)

    def add_connections(self, connections):
        """ Add connections(list of tupules) to graph. """
        for node1, node2 in connections:

    def add(self, node1, node2):
        """ Add node1 and node2 to graph which is initialized with set by default. """

    def get_graph(self):
        return dict(self._graph)

def mex(arr_set):
    mex = 0
    while mex in arr_set:
    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))


  • \$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?


1 Answer 1

    for p in permutations(all_vertices, v):
  • \$1≤n≤10^5\$

Well, \$(10^5)! \approx \left(\frac{10^5}{e}\right)^{10^5} \approx 10^{35657}\$ so it's a waste of time trying to optimise this. The only thing to do is go back to the drawing board and spend a few hours thinking about the mathematics. At best the code you've written will serve to analyse all trees up to a small number of vertices (maybe 7 or 8) to see what patterns you can spot which can help to guide you in the right direction.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.