1 / 3
Jul 2011

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!

  • created

    Jul '11
  • last reply

    Jul '11
  • 2

    replies

  • 269

    views

  • 2

    users

The data structure you use for P is inefficient, if you only want to check for containment.
If you choose a more efficient data structure, you'll get AC.

Indeed. I changed P to a set and Q to a deque, and then the program ran a lot faster. Many thanks to your suggestion on optimization! smiley