HackerRank Solution: BFS Shortest Path

bfs_shortest_path.py

#!/bin/python3
# https://www.hackerrank.com/challenges/bfsshortreach/problem

import math
import os
import random
import re
import sys

from itertools import chain

def bfs(n, m, edges, s):
    distance_map = {node: -1 for node in range(1, n + 1)}
    edge_map = {node: set() for node in range(1, n + 1)}

    # Make a directed graph so we can 'mark' touched nodes by 
    # removing edges *to* them
    for a, b in edges:
        edge_map[a].add(b)
        edge_map[b].add(a)

    # Remove all edges to starting node
    for neighbour in edge_map[s]:
        edge_map[neighbour].remove(s)

    node_queue = iter([s])
    distance_map[s] = 0

    while True:
        try:
            current_node = next(node_queue)
        except StopIteration:
            break
        
        # Invariant: there are no edges *to* current_node
        # Remove edges *from* a node once we have its neighbours
        neighbours = [
            n for n in edge_map.pop(current_node) if n in edge_map
        ]

        # Record distance to neighbours
        for n in neighbours:
            distance_map[n] = distance_map[current_node] + 6
        # Remove all edges *to* neighbours
        for node in edge_map:
            edge_map[node].difference_update(neighbours)

        node_queue = chain(node_queue, neighbours)
    
    distance_map.pop(s)
    return next(zip(*sorted(distance_map.items())))

# ------ HackerRank starter code ------
if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    q = int(input())

    for q_itr in range(q):
        nm = input().split()

        n = int(nm[0])

        m = int(nm[1])

        edges = []

        for _ in range(m):
            edges.append(list(map(int, input().rstrip().split())))

        s = int(input())

        result = bfs(n, m, edges, s)

        fptr.write(' '.join(map(str, result)))
        fptr.write('\n')

    fptr.close()