I'm using iterative dfs with psyco to solve this problem but still kept getting TLE. The following is my code:
from sys import stdin, stdout
def dfs(G, s):
P, Q = [], []
Q.append(s)
while Q:
u = Q.pop()
if u in P:
continue
P.append(u)
for v in G[u]:
Q.append(v)
return P
def main():
num_of_nodes, num_of_edges = map(int, stdin.readline().split())
input_array = stdin.read().splitlines()[:num_of_edges]
if num_of_nodes == 1:
print "YES"
return
if num_of_edges != num_of_nodes-1:
print "NO"
return
graph = {}
for line in input_array:
node, edge = line.split()
if graph.get(node) is None:
graph[node] = []
graph[node].append(edge)
if graph.get(edge) is None:
graph[edge] = []
graph[edge].append(node)
P = dfs(graph, graph.keys()[0])
if len(P) == num_of_nodes:
print "YES"
else:
print "NO"
if __name__ == "__main__":
main()
Is it because the way I create the adjacent graph is too slow? Or are there other bottlenecks with my program?
Thanks a lot!
