felipeblassioli's blog

By felipeblassioli, history, 8 years ago, In English

EDIT: Solved with iterative version and optmizations with 0.982s 28 424 KB. Recursive version throws runtime error at test 31. (non-exit code 0). I think it is segmentation fault

Lesson learned: Python stack frames are heavy, too much recursion is no option if you cant use import resource and get more machine resources.


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

AC Solution (Iterative):

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

    tkns = iter(map(int,stdin.read().split()))
    n = tkns.next()
    m = tkns.next()

    graph = [[] for i in xrange(n+1)]

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

    dp1 = {}
    dp2 = {}
    best = {}
    cache = {}

    root = u
    reverse_topological_order = deque(maxlen=n)
    stk = deque([root],n)
    app = stk.append
    while stk:
        v = stk.pop()
        if len(graph[v]) == 0:
            dp1[v] = dp2[v] = 0
            best[v] = -1
            cache[v] = 0
        else:
            reverse_topological_order.appendleft(v)
            for c in graph[v]:
                graph[c].remove(v)
                app(c)

    for v in reverse_topological_order:
        dp2[v] = sum( cache[c] for c in graph[v] )
        a = dp2[v]
        M = 0
        for c in graph[v]:
            b = 1 + dp2[c] + a - cache[c] 
            if M <= b:
                M = b
                best_c = c

        dp1[v] = M
        if dp2[v] >= dp1[v]:
            best[v] = -1
            cache[v] = dp2[v]
        else:
            best[v] = best_c
            cache[v] = dp1[v]

    w = stdout.write
    stk = deque([root],n)
    w('%d\n' % cache[root])
    ret = []
    app = ret.append
    while stk:
        v = stk.pop()
        if best[v] != -1:
            app('%s %s\n' % (v,best[v]))
        stk.extend(graph[v])
    w(''.join(ret))

main()

Recursive Solution (RTE):

Got RTE probably because of segmentation fault

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

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

    graph = [[] for i in xrange(n+1)]

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

    dp1 = {}
    dp2 = {}
    best = {}
    cache = {}

    def dfs(v):
        if len(graph[v]) == 0:
            dp1[v] = dp2[v] = 0
            best[v] = -1
            cache[v] = 0
            return 

        a = 0
        for c in graph[v]:
            graph[c].remove(v)
            dfs(c)
            a += cache[c]

        dp2[v] = a
        M = 0
        for c in graph[v]:
            b = 1 + dp2[c] + a - cache[c] 
            if M <= b:
                M = b
                best_c = c

        dp1[v] = M
        if dp2[v] >= dp1[v]:
            cache[v] = dp2[v]
            best[v] = -1
        else:
            cache[v] = dp1[v]
            best[v] = best_c

    w = stdout.write
    ret = []
    app = ret.append
    def solve(v):
        for c in graph[v]:
            solve(c)
        if best[v] != -1:
            app('%s %s\n' % (v,best[v]))

    from random import randint
    root = randint(1,n)
    dfs(root)
    w('%d\n' % max(dp1[root],dp2[root]))
    solve(root)
    w(''.join(ret))

main()
  • Vote: I like it
  • +5
  • Vote: I do not like it

| Write comment?
»
8 years ago, # |
Rev. 3   Vote: I like it +8 Vote: I do not like it

Implement the graph by using a list instead of a dictionary, i.e. replace

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

by

graph = [[] for i in xrange(n)]

The lookup for the adjacency list should be faster that way. Note that the name of the vertices then have to be zero-indexed, i.e. their names are 0, 1, ..., n-1.