Dijkstra's Algorithm + an additional cost to start at each node

Question

I am trying to solve the 2009 Canadian Computing Competition Senior #4. The question gives you a map (graph) of different cities, and the cost of transporting something between each city. There are also a few select cities that sell pencils, at a specific cost. You are then given a destination city, and asked to find the cheapest place to buy a pencil from, taking into account the cost of the pencil, and the transportation cost from the pencil seller to the destination city (if the destination city sells pencils, the transportation cost is 0).

# Cities: $$1 \le N \le 5,000$$ # Trade Routes: $$1 \le T \le 5,000,000$$ Route costs: $$0 \le R_{i,j} \le 10,000$$ # Pencils: $$1 \le K \le N$$ Pencil costs: $$0 \le P_i \le 10,000$$

I used Dijkstra's algorithm to find the smallest transportation cost from each city to the destination, and then added it to the cost of a pencil from that city, and then printed out the minimum.

On the judge at DMOJ, I got the first three test cases correct, however, I exceed the time limit of five seconds on the next three test cases. So, I'm interested in knowing if there are any ways I can further optimize my code to save time?

I have already tried using a 2D array instead of a dictionary to store the trade routes, but it makes taking input longer, and I don't how to make Dijkstra's algorithm work better on a 2D array instead of a dictionary.

The reason for slowness could also be due to the act of input() taking too long, which I saw in a discussion about solving the question in C++, but I'm not if that also applies to Python.

# https://dmoj.ca/problem/ccc09s4
from heapq import heappop, heappush

n = int(input())
trade_routes = {i: {} for i in range(1, n + 1)}
pencil_cities = {}

t = int(input())
for _ in range(t):
    data = list(map(int, input().split()))
    trade_routes[data[0]][data[1]] = int(data[2])
    trade_routes[data[1]][data[0]] = int(data[2])

k = int(input())
pencil_costs = {}
for _ in range(k):
    data = list(map(str, input().split()))
    pencil_costs[int(data[0])] = int(data[1])

d = input()


def shortest_path_weighted(graph, start):
    """ (dict, str, str) -> int, list
        Returns the cost and path of the shortest path on a weighted graph
    """
    distances = [100000000 for _ in range(len(graph) + 1)]  # cost from start to each node
    distances[start] = 0

    visited = set()
    visited.add(start)
    queue = [(distances[start], start)]

    while len(queue) > 0:
        (cost, target) = heappop(queue)
        if target not in visited:
            visited.add(target)

        for neighbor in graph[target]:
            if neighbor not in visited:
                temp_distance = distances[target] + graph[target][neighbor]
                if temp_distance < distances[neighbor]:
                    distances[neighbor] = temp_distance
                    heappush(queue, (temp_distance, neighbor))

    return distances


smallest_cost = 99999999
travel_costs = shortest_path_weighted(trade_routes, int(d))

for pencil_cost in pencil_costs:
    if travel_costs[pencil_cost] + pencil_costs[pencil_cost] < smallest_cost:
        smallest_cost = travel_costs[pencil_cost] + pencil_costs[pencil_cost]

print(smallest_cost)

""" Sample Test Case
3
3
1 2 4
2 3 2
1 3 3
3
1 14
2 8
3 3
1
--
6
"""
````

Answer 1

Code Review

While I understand you're looking for performance improvements, this is Code Review, so let's start with seeing if we can improve the code first -- including minor performance improvements -- before we concentrate on algorithmic improvements.

Structure

Code should be organized. With Python, the convention is import statements, then function definitions, then mainline code. You've got a function definition embedded in the mainline code, which is really quite jarring when it comes to reading the code.

Move the function up to just below the import statement.

Only one function? There could be many more! Functions for:

• input number of cities
• inputing the trade route data
• inputing the pencil cost data
• determining smallest manufacting cost + transportation cost

Naming

While n is often used as a stand-in for a counting number, it doesn't tell us what it is the count of. True, in programming contest problem statements, they often give things like n is the number of cities, and t is the number of trade routes, and so on. But look at your code: no where is there even a comment which suggests k is the number of pencil manufacturing cities. How is a reader supposed to determine this, unless they also have the problem statement.

Code should be self documenting, so I would change:

• n → num_cities
• t → num_trade_routes
• k → num_pencil_manufacturers
• d → destination_city

Note that while Python is an "interpreted language", using longer names won't actually slow the program down significantly. There is a "compile" step which happens when the Python source code is read by the interpreter, which ensures the syntax is correct and so on, but it also interns the variable names into tokens, which means variable lookup times will be near constant regardless of variable name length.

Heterogeneous variables

    data = list(map(int, input().split()))
    trade_routes[data[0]][data[1]] = int(data[2])
    trade_routes[data[1]][data[0]] = int(data[2])

Continuing with naming, what is data? What is data[0]? What is data[2]? The code doesn't give us a clue!

data is a variable which holds a list of ... data items. Specifically, it holds two cities and a route cost, so it needs to be called ... uhm, two_cities_and_route_cost? Yuck!

The problem here is there are 2 different types of data in that list -- cities and cost -- so there isn't a good name for that. This means they shouldn't be stored inside one variable.

The code has a few other problems: there is no check that data actually contains 3 values. It could hold 5, and two would be silently forgotten. It could hold only one or two, and the next statement would actually generate an IndexOutOfBound exception, but that is too late; the error would be flagged on the wrong statement, and the programmer is forced to trace through the code to discover where the problem actually started.

The .split() function has an optional maxsplit parameter, which could be used to ensure you don't end up with more than 3 values. While that helps, it doesn't go far enough.

For this issue Python has "tuple unpacking". Assigning a list of data into a tuple ensures that the list has exactly the correct number of items. It also gives names to the values, solving the early point.

Finally, the int(...) in int(data[2]) is redundant, as map(int, ...) has already converted the value to an int.

    city1, city2, cost = map(int, input().split())
    trade_routes[city1][city2] = cost
    trade_routes[city2][city1] = cost

This code is much clearer, and safer. input() is .split() into an unknown number of terms. Then map(int, ...) returns an iterable which will convert each term in succession into an integer ... but nothing has happened yet. Previously you used list(...) to force the iteration. Now, we simply assign the iterable to a tuple, which causes tuple-unpacking. The tuple has three elements, so the iterable MUST resolve itself into 3 values for the assignment to succeed. If more or less than three values are given, this statement itself raises the exception ... not usage of a data variable at some (much) later point in the code.

The heterogeneous data variable has been eliminated AND we've given names to each of the values. Win-win!

The same should be applied to the pencil costs input.

Incorrect Documentation

The only thing worse than no documentation is incorrect documentation.

def shortest_path_weighted(graph, start):
    """ (dict, str, str) -> int, list
        Returns the cost and path of the shortest path on a weighted graph
    """
    ...
    return distances

The """docstring""" seems to suggest the function takes 3 parameters (a dict and two strings) and returns an integer (the cost?) and a list (the path through the graph?). The function signature only has 2 parameters, and when it is called it is a passed trade_routes, int(d) ... the second argument which is clearly an int; there is no string in sight!

Use type hints, and use a linter which actually checks the type hints.

Initialization

    visited = set()
    visited.add(start)

Why two statements, when one will do?

    visited = {start}

Magic numbers

What is 100000000, and where did it come from? How about 99999999?

Use named constants where possible.

MAX_CITIES = 5_000
MAX_ROUTE_SEGMENT_COST = 10_000
MAX_PENCIL_COST = 10_000
MAX_TRANSPORT_COST = MAX_CITIES * MAX_ROUTE_SEGMENT_COST
MAX_TOTAL_COST = MAX_PENCIL_COST + MAX_TRANSPORT_COST

Or don't worry about that exact maximum value, and simply use math.inf.

Busy Work

        if target not in visited:
            visited.add(target)

Why two statements? visited is a set; adding target to visited if it is already in the set is a no-op. You could just write:

        visited.add(target)

Bug/inefficiency?

Or are you supposed skip the for neighbour in graph[target]: loop if target is visited? In which case, you want to keep the test, but move the loop into the body of the if statement.

Algorithmic Improvements

Consider the test case: $$N = 5000$$ $$C_{i,j} = j - i, 1 \le i \lt j \le 5000$$ $$K = 1$$ $$P_2 = 10$$ $$D = 1$$

5000 cities, all in a straight line, with fully connected trade routes, and pencils only manufactured in the second city.

Question: do you really need to calculate the minimum distance to every other city? No. As soon as you've expanded your minimum distance to include the last pencil manufacturing city (in this case, the only one), you can stop computing the transportation costs ... in this case ignoring the remaining 4998 cities, and 4998*4997/2 trade routes, which would be a considerable time savings.

Let's add another pencil manufacturer, \$P_{4000} = 4\$. Well, now we have to continue exploring a bit further. How much? After \$C_{1,8} = 7\$, a pencil costing 4 won't have a net cost lower than the 11 from City #2, so we don't have to explore until we reach all of the pencil cities ... just until there are no pencils cheaper than the cheapest total cost less the transportation distance we've explored to.

That's a complicated stopping condition, but if it terminates the search before exploring the entire network, it should be worthwhile.

But we can make things even simpler by changing the problem slightly:

• Only one pencil manufacturing city: city #0
• Manufacturing cost is 0
• \$R_{0,i} = P_i, 1 \le i \le K\$

We've added a route with a cost equal to the manufacture cost of the pencil at each pencil manufacturing city. Now the problem is merely the shortest weighted path from city #0 to the destination city: $$R_{0,D}$$ and early termination should be trivial.