Search code examples
pythonpython-3.xgraph-theoryshortest-pathdijkstra

In a directed, weighted graph, efficiently find the cost of the shortest path between two nodes A and B that traverses through a node X


I'm a beginner and still grasping the concept of pathing algorithms. Our professor has given us this problem which solution will be evaluated by an online judge:

I am given the following:

  • The number of nodes L (each node is named as a number, 1-indexed)
    • 2 <= L <= 1000
  • The directed edges from one node to another R, as well as the cost C of the edge
    • 0 <= R <= L * sqrt(L)
    • 1 <= C <= 100
    • The number of edges is given alongside the number of nodes to easily gather the directed edges.
    • There are no negative costs.
  • the source node A
  • the destination node B
  • the traversed node X

I need to find the cost of the shortest path from A to X to B.

My code runs as follows:

  1. Get the input.
  2. Generate the adjacency list of the graph.
  3. Get A, X and B.
  4. Get the cost ax of the shortest path between A and X.
  5. Get the cost xb of the shortest path between X and B.
  6. The cost axb of the shortest path from A to X to B is ax + xb.

The judge evaluates this solution as exceeding the time limit of 1 second. Is there a way to improve this and make it more efficient?

  • I've considered making X the source node, but the graph is directed, and undirecting it will yield the wrong results.

  • I've considered checking the existence of the path between A and X, then X and B, but it seems to have negligible effect on the performance.

  • I've also considered "clamping" the graph so that it only includes paths between A and B that goes through X, but I don't know how to do that in the most efficient manner.

  • I've tried applying this idea but it just made it run longer.

  • My professor says A* would be overkill for the problems he gave us, but if I figure out a heuristic for this maybe I'll consider using it.

  • I tried the Floyd-Warshall Algorithm, but it made the code consume more time and memory - but something tells me this algorithm can be optimized:

def shortestPath(dist, i, j, k): # My first try in applying the Floyd-Warshall Algorithm

    for k in range(len(dist)):
        for i in range(len(dist)):
            for j in range(len(dist)):
                if dist[i][j] > dist[i][k] + dist[k][j]:
                    dist[i][j] = dist[i][k] + dist[k][j]

# Graph construction and algorithm application:
    .
    .
    .
    # generate adjacency list of the graph
    graph = [[math.inf for node_x in range(nodes)] for node_y in range(nodes)]
    next = [[None for node_x in range(nodes)] for node_y in range(nodes)]
    for node in range(nodes):
        graph[node][node] = 0

    for edge in range(edges):
        node_a, node_b, cost = input().split()
        node_a = int(node_a) - 1
        node_b = int(node_b) - 1
        cost = eval(cost)
        graph[node_a][node_b] = cost
        next[node_a][node_b] = node_b

    i, k, j = input().split()  # get source, traversed and destination node
    i = int(i) - 1
    j = int(j) - 1
    k = int(k) - 1
    shortestPath(graph, i, j, k)
    ikj = graph[i][j]

    if ikj == math.inf:  # path nonexistent if either ax or xb is infinite
        out.append("path nonexistent")
        continue
    else:
        out.append(ikj)

This is my code:

import heapq
import math


# tools start here

def dijkstra(G, S):
    pq = []
    costs = {v: math.inf for v in G}
    costs[S] = 0

    heapq.heappush(pq, (0, S))
    while pq:
        d_u, u = heapq.heappop(pq)
        for e in G[u]:
            v, w = e
            d_v = costs[v]
            if d_v > d_u + w:  # relaxation operation
                costs[v] = d_u + w
                heapq.heappush(pq, (d_u + w, v))
    return costs


# tools end here

t = int(input())  # get number of test cases
out = []  # initialize output array
for _ in range(t): # for each test case

    # get input
    nodes, edges = input().split()
    nodes = int(nodes)
    edges = int(edges)

    # generate adjacency list of the graph
    graph = {}
    for node in range(1, nodes + 1):
        graph[str(node)] = []

    for edge in range(edges):
        node_a, node_b, cost = input().split()
        cost = eval(cost)
        graph[node_a].append((node_b, cost,))

    a, x, b = input().split()  # get source, traversed and destination node

    ax = dijkstra(graph, a)[x]  # get shortest path cost from a to x
    xb = dijkstra(graph, x)[b]  # get shortest path cost from x to b

    axb = ax + xb  # add the path costs

    if axb == math.inf:  # path nonexistent if either ax or xb is infinite
        out.append("path nonexistent")
        continue
    else:
        out.append(axb)

[print(_) for _ in out]  # print output array

Here is the sample input:

2
3 2
1 2 1
2 3 1
1 2 3
3 1
1 2 1
1 2 3

Here is the sample output:

2
path nonexistent

Solution

  • I think your solution is already the fastest.

    If you want to squeeze out more runtime performance from your code, I recommend reincorporating the check for path existence between A and X, then X and B.

    I also recommend replacing your eval() with int(), since it would hypothetically be slower to check if it's anything other than an int, considering your problem's constraints.

    Finally, this is likely an online judge that accepts code that prints the output as it goes, so output the answer immediately after every test case.

    Hope that helps!