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()
Iterative Solution:
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:
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()