Cheapest Flights Within K Stops

Medium

The problem

There are n cities numbered 0 to n - 1. Each entry [from, to, price] in flights is a one-way flight. Return the cheapest price to travel from src to dst using at most k stops (cities in between), or -1 if that is impossible.

Example 1
Input: n = 4, flights = [[0, 1, 100], [1, 2, 100], [2, 0, 100], [1, 3, 600], [2, 3, 200]], src = 0, dst = 3, k = 1
Output: 700
With one stop, the only way is 0 to 1 to 3 for 100 + 600.
Example 2
Input: n = 4, flights = [[0, 1, 100], [1, 2, 100], [2, 0, 100], [1, 3, 600], [2, 3, 200]], src = 0, dst = 3, k = 2
Output: 400
With two stops you can go 0 to 1 to 2 to 3 for 100 + 100 + 200.
Example 3
Input: n = 3, flights = [[0, 1, 100], [1, 2, 100], [0, 2, 500]], src = 0, dst = 2, k = 0
Output: 500
With no stops, only the direct flight is allowed.

Limits

1 ≤ n ≤ 100
0 ≤ k < n
0 ≤ len(flights) ≤ 5,000
1 ≤ price ≤ 10,000
src and dst are different

Write it in Go. Try for about 20 minutes on paper first, then open one hint at a time.

Try it here

Write Go. Common packages like fmt and sort are imported for you. Keep the function name and inputs the same.

Hints, one at a time

Nudge

You have at most k+1 flights. What if you relax all edges exactly one round at a time?

The idea

Bellman-Ford limited to k+1 rounds, each round reading from a copy of the previous round's distances. (Dijkstra with state (node, stops) also works.)

Target: O(k · E) time

Go function shape

func findCheapestPrice(n int, flights [][]int, src int, dst int, k int) int

Reference solution

Tested with go test. Try it yourself first, then compare.

// FindCheapestPrice: cheapest src→dst price with at most k stops.
// Bellman-Ford limited to k+1 rounds. Each round reads the PREVIOUS round's prices, so one round = one more flight.
func FindCheapestPrice(n int, flights [][]int, src, dst, k int) int {
	prices := make([]int, n)
	for i := range prices {
		prices[i] = math.MaxInt
	}
	prices[src] = 0
	for range k + 1 {
		next := append([]int(nil), prices...)
		for _, f := range flights {
			from, to, cost := f[0], f[1], f[2]
			if prices[from] != math.MaxInt && prices[from]+cost < next[to] {
				next[to] = prices[from] + cost
			}
		}
		prices = next
	}
	if prices[dst] == math.MaxInt {
		return -1
	}
	return prices[dst]
}

Read the lesson: Advanced Graphs · Dijkstra & Spanning Trees