Module assignment: City route planner in Python

The brief

A regional transport office keeps a list of towns and the roads between them. Each road has a travel time in minutes, and every road can be used in both directions. Staff at the office type questions into a terminal, and your program answers them. Four kinds of question come in:

• ROUTE s t: what is the fastest travel time from town s to town t?
• HOPS s t: what is the fewest number of roads needed to go from s to t, ignoring the minutes?
• GROUPS: how many separate groups of towns are there, where a group is a set of towns that can all reach each other?
• PLAN: what is the smallest total number of minutes of road that would connect every town, if the office built only the roads it needs?

This assignment uses everything in the module. ROUTE needs the weighted shortest paths from lesson 4, which use a heap. A heap is a structure that always gives back its smallest item first, and Python provides one in the heapq module. HOPS needs the breadth-first search from lesson 1. GROUPS needs union-find or a depth-first count from lessons 2 and 5. PLAN needs Kruskal's minimum spanning tree from lesson 5. Some of the answers can come from more than one method. Pick the one that fits the question, and say why in a comment.

You will write one Python program that reads from standard input and writes to standard output. Standard input is the text that goes into the program, and standard output is the text it prints. The Input box in this course runs your program with the text you type, so make sure the output matches the examples exactly.

The input format

The first line holds two numbers: n, the number of towns, and m, the number of roads. Towns are numbered from 0 to n - 1.

The next m lines describe the roads. Each line has three numbers, u v w, meaning there is a road between town u and town v that takes w minutes. Roads work both ways. A pair of towns may be joined by more than one road, and a town never has a road to itself.

After the roads comes a line with one number, q, the number of queries. The next q lines each hold one query, in one of the four forms above. A query is one word, followed by two town numbers for ROUTE and HOPS, and nothing else for GROUPS and PLAN.

The limits are: 1 <= n <= 2000, 0 <= m <= 10000, 1 <= w <= 1000000 and 0 <= q <= 500. The longest possible answer is 1999000000 minutes (1999 roads of 1000000 each), which only just fits in a 32-bit integer. Python's int has no fixed size, so the minutes need no special type. A Java program would use long for the same job.

The output format

Print one line for each query, in the same order as the queries. The format of each answer is set out below.

• ROUTE prints the fastest number of minutes. If the destination cannot be reached, print UNREACHABLE. If the start and the destination are the same town, the answer is 0.
• HOPS prints the fewest number of roads. If the destination cannot be reached, print UNREACHABLE. The answer for the same town is 0.
• GROUPS prints the number of groups. A town with no roads is a group on its own.
• PLAN prints the total minutes of the cheapest set of roads that joins every town. If some towns cannot be joined by any roads, print NO PLAN. A single town needs no roads, so the answer for n = 1 is 0.

Example one

Here is a small network with seven towns, numbered 0 to 6. Town 6 has no roads at all. The following input has eight queries.

7 7
0 1 4
0 2 2
1 2 1
1 3 5
2 3 8
3 4 3
4 5 2
8
ROUTE 0 5
HOPS 0 5
GROUPS
PLAN
ROUTE 0 6
HOPS 2 2
ROUTE 5 0
GROUPS

The expected output is:

13
4
2
NO PLAN
UNREACHABLE
0
13
2

Check the first answer by hand, because it is the one most likely to go wrong. The fastest way from town 0 to town 5 is 0 to 2 (2 minutes), then 2 to 1 (1 minute), then 1 to 3 (5 minutes), then 3 to 4 (3 minutes), then 4 to 5 (2 minutes). That adds up to 13. Going 0, 1, 3 costs 4 + 5 = 9 minutes, while 0, 2, 1, 3 costs 2 + 1 + 5 = 8, so the cheaper way to town 3 goes through town 2. The answer for HOPS 0 5 is 4, which is the route 0, 1, 3, 4, 5. The route 0, 2, 3, 4, 5 also has 4 roads, so either route gives the same answer.

Notice that HOPS and ROUTE give different answers for the same pair of towns. HOPS does not care about minutes, and ROUTE does. Your program must keep the two apart.

Example two

The second example has four towns, a different set of roads, and a network that can be joined into one piece.

4 5
0 1 3
1 2 3
2 3 1
0 3 7
1 3 2
3
PLAN
GROUPS
ROUTE 0 2

The expected output is:

6
1
6

For PLAN, the cheapest three roads are 2-3 (1 minute), 1-3 (2 minutes) and 0-1 (3 minutes), which join all four towns for 6 minutes in total. The road 1-2 also costs 3, but taking it would form a loop with 1-3 and 2-3, so it is skipped. For ROUTE 0 2, the fastest path is 0, 1, 3, 2 at 3 + 2 + 1 = 6 minutes. The direct route 0, 3, 2 costs 7 + 1 = 8, and the route 0, 1, 2 costs 3 + 3 = 6, so two routes tie at 6. Either gives the same answer.

Testing your program

Examples are only a start. Before you submit, write at least three inputs of your own. Include these cases:

• A single town and no roads, with the queries GROUPS, PLAN and ROUTE 0 0. The answers are 1, 0 and 0.
• Two towns joined by one road, with a query in each direction. The answers are the same both ways, because roads work in both directions.
• A network where the fewest-roads route is not the fastest route. Make one road very slow and build a longer route that is quicker. Check that ROUTE and HOPS disagree, as they should.
• A network with a town that is not reachable from the others. Check that ROUTE prints UNREACHABLE, that GROUPS counts the isolated town, and that PLAN prints NO PLAN.

For example, the fewest-roads test can be this input. The road 0-1 is slow, so the answers differ:

4 3
0 1 100
1 2 1
0 2 1
2
ROUTE 0 1
HOPS 0 1

The output is 2 for ROUTE (0, 2, 1 at 1 + 1 minutes) and 1 for HOPS (the direct road).

For each of your own tests, work out the answer on paper before you run the program. A program that agrees with a wrong answer you calculated by hand has a bug in the program, the paper, or both, so look again at the difference.

Requirements checklist

Copy this list into your notes and tick each item when you have checked it.

• The program reads n, m, the m roads and q from standard input, in that order, with any spacing between numbers.
• ROUTE uses a shortest-path method that handles weights, and prints the minutes as a whole number.
• ROUTE prints UNREACHABLE when the destination cannot be reached.
• HOPS ignores the minutes and counts roads, using a breadth-first search.
• HOPS prints UNREACHABLE when the destination cannot be reached.
• GROUPS counts isolated towns as groups of their own.
• PLAN uses Kruskal's algorithm or Prim's algorithm, and prints NO PLAN when the towns cannot all be joined.
• PLAN prints 0 for a single town with no roads.
• Every query prints exactly one line, in the same order as the queries.
• The program passes both examples in this lesson, with output matching exactly, including the line breaks.
• The program passes at least three inputs of your own, checked against hand calculations.
• The program is cut into functions, one for each kind of query, with a short comment on each function for why that method was chosen.
• The program runs fast enough for the largest limits: 2000 towns, 10000 roads and 500 queries should finish in a few seconds. The reference solution below took about 2 seconds on one generated input at these limits. Your time will differ.

Grading rubric

Your program is graded on the answers it prints across a set of hidden test inputs. The grade is a percentage of the tests it gets right, with a few marks for clean code. The three levels below describe what each grade looks like in practice.

A solution that gets the examples right but fails on the largest limits is at 80% at most, because the hidden tests include large inputs. A solution that uses the wrong algorithm for a query, for example breadth-first search for ROUTE, is at 60% at most, even when its answers happen to be right on small inputs. The reason is that the method, not just the answer, is what the module teaches.

Hints

Try the work yourself first. Each hint below opens one step at a time. Read one hint, try again, and open the next only if you are stuck.

Hint 1: start with the reading code.

Show the hint

Read all of the input with sys.stdin.read().split(). This gives a list of words. Turn the list into an iterator with tokens = iter(words), and take the next word with next(tokens). Build the adjacency list as adj = [[] for _ in range(n)], and add each road to both towns' lists as a (neighbour, weight) pair. Store each road as a (u, v, w) tuple in a separate list too, because PLAN needs the list of all roads. Read q and then loop q times, reading the query word with next(tokens) and the numbers that follow it with int(next(tokens)). Print nothing until the reading loop works.

Hint 2: ROUTE is Dijkstra.

Show the hint

Copy the Dijkstra program from lesson 4. Use math.inf as the distance for every town at the start. Put the start in a heap with heapq.heappush, storing tuples of the form (distance, town). A heap pops the smallest tuple first, so heapq.heappop always gives the closest unfinished town. Return the whole list of distances, so that the answer for any destination is available. Print UNREACHABLE when the value is still math.inf.

Hint 3: HOPS is BFS with the -1 marker.

Show the hint

Copy the breadth-first search from lesson 1. The list of distances starts at -1, and you set it to 0 for the start. Use collections.deque as the queue, with append to add at the back and popleft to take from the front. There are no weights, so each road adds exactly 1. The answer for an unreachable town is still -1, and that is the signal to print UNREACHABLE.

Hint 4: GROUPS and PLAN share union-find.

Show the hint

Write a find function with path halving, as in lesson 5. For GROUPS, start the count at n, and go through the roads, decreasing the count each time a road joins two different groups. For PLAN, sorted(edges, key=...) returns a sorted copy of the roads, so you can sort by weight with key=lambda e: e[2]. Run the same union loop on that copy, adding the weight of each road that joins two groups. Print NO PLAN when the number of roads taken is not n - 1.

Hint 5: reset the union-find before each query.

Show the hint

Both GROUPS and PLAN change the parent list. If the two functions shared one list, the second answer would start from the first answer's groups and be wrong. Create parent = list(range(n)) at the start of each function, not once at the top of the file. Each function then has its own list, and the reset happens without extra code.

Reference solution

Try your own version before you open this. The solution below is a complete program. It passes both examples in this lesson. Read it after you have a working version, and compare the choices, especially the way it keeps the shortest-path and fewest-roads answers apart.

Show the reference solution

import heapq
import sys
from collections import deque
from math import inf
 
 
def find(parent, x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x
 
 
# ROUTE: Dijkstra, because roads have weights and we need the smallest total.
def fastest(adj, n, start):
    dist = [inf] * n
    dist[start] = 0
    heap = [(0, start)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist[u]:
            continue
        for v, w in adj[u]:
            candidate = d + w
            if candidate < dist[v]:
                dist[v] = candidate
                heapq.heappush(heap, (candidate, v))
    return dist
 
 
# HOPS: BFS, because every road counts as one hop whatever its minutes.
def hops(adj, n, start):
    dist = [-1] * n
    dist[start] = 0
    queue = deque([start])
    while queue:
        u = queue.popleft()
        for v, _ in adj[u]:
            if dist[v] == -1:
                dist[v] = dist[u] + 1
                queue.append(v)
    return dist
 
 
# GROUPS: union-find, starting from n separate towns and merging along roads.
def groups(n, edges):
    parent = list(range(n))
    count = n
    for u, v, _ in edges:
        ru = find(parent, u)
        rv = find(parent, v)
        if ru != rv:
            parent[ru] = rv
            count -= 1
    return count
 
 
# PLAN: Kruskal, because the cheapest set of roads that joins everything is a minimum spanning tree.
def plan(n, edges):
    parent = list(range(n))
    total = 0
    taken = 0
    for u, v, w in sorted(edges, key=lambda e: e[2]):
        ru = find(parent, u)
        rv = find(parent, v)
        if ru != rv:
            parent[ru] = rv
            total += w
            taken += 1
    return total if taken == n - 1 else None
 
 
def main():
    tokens = iter(sys.stdin.read().split())
    n = int(next(tokens))
    m = int(next(tokens))
    adj = [[] for _ in range(n)]
    edges = []
    for _ in range(m):
        u = int(next(tokens))
        v = int(next(tokens))
        w = int(next(tokens))
        edges.append((u, v, w))
        adj[u].append((v, w))
        adj[v].append((u, w))
 
    q = int(next(tokens))
    lines = []
    for _ in range(q):
        kind = next(tokens)
        if kind == "ROUTE":
            s = int(next(tokens))
            t = int(next(tokens))
            d = fastest(adj, n, s)[t]
            lines.append("UNREACHABLE" if d == inf else str(d))
        elif kind == "HOPS":
            s = int(next(tokens))
            t = int(next(tokens))
            h = hops(adj, n, s)[t]
            lines.append("UNREACHABLE" if h == -1 else str(h))
        elif kind == "GROUPS":
            lines.append(str(groups(n, edges)))
        elif kind == "PLAN":
            total = plan(n, edges)
            lines.append("NO PLAN" if total is None else str(total))
    sys.stdout.write("".join(line + "\n" for line in lines))
 
 
if __name__ == "__main__":
    main()

Why the reference solution is built this way

Each query gets its own function, and the function name says what it does. That makes the program easy to check against the requirements. ROUTE and HOPS both start from a town and return a whole list of answers, so the same reading of the roads can serve every query of that kind. GROUPS and PLAN each create their own parent list at the start, which is the fix described in hint 5.

The edges list is kept separate from the adjacency list on purpose. Kruskal needs the list of roads in a form it can sort. sorted returns a new list and leaves edges in its original order, so the other queries are not affected. A Java version needs to copy the array first to get the same result. In Python, sorted makes the copy.

Dijkstra starts every distance at math.inf, which is a number larger than any real distance. Python compares an int with inf directly, so candidate < dist[v] needs no special case for towns that have not been reached yet.

The output is collected in a list of lines and written once at the end. sys.stdout.write with the joined lines prints exactly one line per query. It also prints nothing when q is 0, where print("") would print one blank line.

The last two lines use if __name__ == "__main__":. This runs main only when the file is run directly. If another program imports this file, the functions can be used without reading any input.

Submitting your work

Run your program on both examples in the Input box, and compare the output with the expected output line by line. A common slip is an extra blank line at the end, or a missing line when a query prints nothing. Each query must print exactly one line.

Then run your three own tests. For each, write down which answer you expected, and say which lines of your program produced it. If a line surprises you, go back to the function for that query and trace it by hand with a small example.

When you are satisfied, submit the program as the answer to this assignment, together with a short note of two or three sentences. The note should name the algorithm you used for each of the four queries, and explain in one sentence why ROUTE and HOPS give different answers for the same pair of towns. This note is part of the grade, because the rubric asks for the method and not only the result.