Longest path in a undirected and unweighted tree - in python

Question

I am solving a problem in which I need to calculate the diameter of the tree.I know how to calculate that using 2 bfs first to find the farthest node then do second bfs using the node we found from the first one.

But I am having difficulty to implement a very simple step - making a adjacency list (dict in case of python) from the input I have written a code to this but its not tidy and not the best can someone tell how to do this efficently

Input

The first line of the input file contains one integer N --- number of nodes in the tree (0 < N <= 10000). Next N-1 lines contain N-1 edges of that tree --- Each line contains a pair (u, v) means there is an edge between node u and node v (1 <= u, v <= N).

Example :

8

1 2

1 3

2 4

2 5

3 7

4 6

7 8

My code is :

def makedic(m):
    d = {}
    for i in range(m):
        o, p = map(int, raw_input().split())
        if o not in d and p not in d:
             d[p] = []
             d[o] = [p]
        elif o not in d and p in d:
             d[o] = []
             d[p].append(o)
        elif p not in d and o in d:
            d[p] = []
            d[o].append(p)
        elif o in d:
            d[o].append(p)
           # d[p].append(o)
        elif p in d:
            d[p].append(o)
    return d

Here is how I implemented bfs:

def bfs(g,s):
    parent={s:None}
    level={s:0}
    frontier=[s]
    ctr=1
    while frontier:
        next=[]
        for i in frontier:
            for j in g[i]:
                if j not in parent:
                    parent[j]=i
                    level[j]=ctr
                    next.append(j)
        frontier=next
        ctr+=1
     return level,parent

Answer 1

You are using a 0 instead of o in makedic in first elif condition. Also made the correction for undirected graph.

def makedic(m):
    d = {}
    for i in range(m):
        o, p = map(int, raw_input().split())
        if o not in d and p not in d:
             d[p] = [o]
             d[o] = [p]
        elif o not in d and p in d:
             d[o] = [p]
             d[p].append(o)
        elif p not in d and o in d:
            d[p] = [o]
            d[o].append(p)
        elif o in d:
            d[o].append(p)
            d[p].append(o)
    return d

Answer 2

There are unnecessary checks in your code. For each A - B edge you just have to put B in A 's adjacency list and A in B 's adjacency list:

d = {}
for i in range(m):
    u,v = map(int, raw_input().split())
    if u in d:
        d[u].append(v)
    else:
        d[u] = [v]
    if v in d:
        d[v].append(u)
    else:
        d[v] = [u]   

According to the question every node has an index between 1 and N so you can use this fact and pre-populate the dict with empty lists. This way you don't have to check whether a key is in the dict or not. Also make the code a little bit shorter:

N = input()
d = { i:[] for i in range(1, N+1) }
for i in range(N):
    u,v = map(int, raw_input().split())
    d[u].append(v)
    d[v].append(u)