Depth-first_search.py

"""
There are N countries {1,2,3,4....N} and N-1 roads(i.e depicting a tree)

Bishu lives in the Country 1 so this can be considered as the root of the tree.

Now there are Q girls who live in various countries (not equal to 1) .

All of them want to propose Bishu. But Bishu has some condition.

He will accept the proposal of the girl who lives at minimum distance from his country.

Now the distance between two countries is the number of roads between them.

If two or more girls are at the same minimum distance then he will accept the proposal of the girl who lives in a country with minimum id.

No two girls are at same country.

Input: First line consists of N,i.e number of countries Next N-1 lines follow the type u v which denotes there is a road between u and v. Next line consists of Q Next Q lines consists of x where the girls live.

Output: Print the id of the country of the girl who will be accepted.

contraints: 2<=N<=1000 1<=u,v<=N 1<=Q<=(N-1)
"""


from collections import defaultdict

graph = defaultdict(list)

# input n countries
n = int(input().strip())
for _ in range(n - 1):
    x, y = map(int, input().strip().split())
    graph[x].append(y)
    graph[y].append(x)

# the lines denoting distance x, where the girls live
m = int(input().strip())
girls_city = []
for _ in range(m):
    girls_city.append(int(input().strip()))

girls_city = sorted(girls_city)


def dfs(start):
    stack = [start]
    visited = set()
    visited.add(start)
    _distance = {start: 0}
    while stack:
        node = stack.pop()
        for connected_node in graph[node]:
            if connected_node not in visited:
                visited.add(connected_node)
                stack.append(connected_node)
                _distance[connected_node] = _distance[node] + 1

    return _distance


distance = dfs(1)
distance = [(key, value) for key, value in distance.items()]
distance = sorted(distance, key=lambda node: node[1])
min_distance = None
nearest_city = None
for key, value in distance:
    if not value or key not in girls_city:
        continue

    if nearest_city is None:
        min_distance = value
        nearest_city = key
        continue

    if min_distance == value and key < nearest_city:
        nearest_city = key
        continue

    if value > min_distance:
        break

print("\n"+str(nearest_city))

Sample input format:
6
1 2
1 3
1 4
2 5
2 6
4
5
6
3
4

This should give an answer of 3