How to optmize python solution? The problem is Maximal Matching in a Tree.

Revision en2, by felipeblassioli, 2015-10-18 10:03:36

I am trying to pass a timus problem in python, but I still get TLE in test 23. I tried looots of stuff, so do you guys have any idea?

The input is quite big, I've passed some problems in python with big (10^4, 10^5) inputs and its always quite a pain. Input has a maximum of 100000 lines. And it a graph that is a tree.

Example:

Maybe I should read the input in other way? Maybe I could print the dp solution in a better way? (is solve() really needed?)

Here's the code:

def main():
    from sys import stdin, setrecursionlimit, stdout
    from itertools import islice, izip
    from collections import deque

    setrecursionlimit(10**5+5)
    tkns = iter(stdin.read().split())
    n,m = int(tkns.next()), int(tkns.next()) 

    graph = {str(i): [] for i in xrange(1,n+1)}

    for u,v in izip(tkns, tkns):
        graph[u].append(v)
        graph[v].append(u)

    dp1 = {}
    dp2 = {}
    parent = {}

    visited = set()
    def dfs(v):
        visited.add(v)
        if len(graph[v]) == 0:
            dp1[v] = dp2[v] = 0
            return 

        a = 0
        tmp = deque()
        for c in graph[v]:
            if c not in visited:
                parent[c] = v
                dfs(c)
                a += max(dp1[c],dp2[c])
                tmp.append( 1 + dp2[c] - max(dp1[c],dp2[c]) )

        dp2[v] = a
        if len(tmp) > 0:
            dp1[v] = max( t + a for t in tmp )
        else:
            dp1[v] = 0

    w = stdout.write
    def solve(v):
        a = dp1[v]
        b = dp2[v]
        if b >= a:
            for c in graph[v]:
                if c != parent[v]:
                    solve(c)
        else:
            best = 0
            for c in graph[v]:
                if c != parent[v]:
                    t = dp2[c]
                    x = 1 + t + b - max(t,dp1[c])
                    if x >= best:
                        best = x
                        best_c = c
                    solve(c)
            w('%s %s\n' % (v,best_c))

    root = u
    parent[root] = None
    dfs(root)
    w('%d\n' % max(dp1[root],dp2[root]))
    solve(root)

main()
Tags dp on a tree, matching, graphs, tree, python, python 2, python in competitions, optimization

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en5 English felipeblassioli 2015-10-19 00:09:20 99
en4 English felipeblassioli 2015-10-19 00:06:30 2 Tiny change: 'sources.\n___\n\nI' -> 'sources.\n\n___\n\nI'
en3 English felipeblassioli 2015-10-18 22:58:32 3491 Posted solution
en2 English felipeblassioli 2015-10-18 10:03:36 60
en1 English felipeblassioli 2015-10-18 10:01:38 2320 Initial revision (published)