787. Cheapest Flights Within K Stops

Medium

Problem:

Calculate the cheapest price from the starting point to the destination, but return the price that arrives within k stopovers. If a route does not exist, return -1.

Input: n = 3, flights = [[0,1,100],[1,2,100],[0,2,500]], src = 0, dst = 2, k = 0
Output: 500
Explanation:
The graph is shown above.
The optimal path with no stops from city 0 to 2 is marked in red and has cost 500.

https://leetcode.com/problems/cheapest-flights-within-k-stops/

Solution:

Time Limit Exceeded

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph = collections.defaultdict(list)
        for u, v, w in flights:
            graph[u].append((v, w))

        # Queue = (price, stop, src)
        q = [(0, 0, src)]
        min_price = float('inf')

        while q:
            price, stop, node = heapq.heappop(q)

            # If we've reached the destination, update the minimum price
            if node == dst:
                min_price = min(min_price, price)

            # If we haven't exceeded the allowed number of stops 
            # and the current price is less than the minimum price
            if stop <= k and price <= min_price:
                for v, w in graph[node]:
                    heapq.heappush(q, (price + w, stop + 1, v))

        return min_price if min_price != float('inf') else -1

Logic error

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph = collections.defaultdict(list)
        for u, v, w in flights:
            graph[u].append((v, w))

        # Queue = (price, stop, src)
        q = [(0, 0, src)]
        dist = collections.defaultdict(list)
        result = collections.defaultdict(list)
        while q:
            stop, price, node = heapq.heappop(q)
            if node in dist:
                continue

            dist[node] = (stop, price)
            for v, w in graph[node]:
                if v == dst:
                    result[stop + 1].append(price + w)
                    print("dst: ", result)

                heapq.heappush(q, (stop + 1, price + w,  v))

        cheapest = set()
        for stop in list(result):
            if stop <= k + 1:
                cheapest.add(min(result[stop]))
                
        return min(cheapest) if cheapest else -1

The problem I was facing is due to the way I was handling the dist dictionary. I was marking a node as visited once I've popped it from the priority queue, but this approach doesn't account for the possibility of reaching the same node through a different path with fewer stops or a lower cost.

Let's address the issues:

Instead of using the dist dictionary to track visited nodes, use it to track the minimum cost to reach each node. This way, you can prune paths that have a higher cost than a previously discovered path to the same node.
The result dictionary is not necessary. You can track the minimum cost to reach the destination directly.
The logic for checking the number of stops (k) needs to be incorporated into the while loop to prune paths that exceed the allowed number of stops.

Modified, but still get "time limit exceeded"

Implement it in the same way as the previous Dijkstra's algorithm, but when adding to the priority queue for arriving within k stops, only add the path if it's within k, so that paths exceeding k are no longer explored.

Since there's no need to store the entire distance anymore, the dist dictionary is removed.
Only the shortest path to the destination needs to be calculated, so there's no need to check the total number of paths either.

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph = collections.defaultdict(list)
        for u, v, w in flights:
            graph[u].append((v, w))

        # Queue: [(price, node, k)]
        q = [(0, src, k)]

        while q:
            price, node, count = heapq.heappop(q)
            if node == dst:
                return price
            if count >= 0:
                for v, w in graph[node]:
                    heapq.heappush(q, (price + w, v, count - 1))
        
        return -1

This code uses a breadth-first search (BFS) strategy with a priority queue to ensure that the path with the smallest cost is explored first. However, the algorithm can end up processing nodes multiple times, leading to an increase in the number of operations, especially when the graph is dense.

Optimization: Dijkstra's Algorithm

A way to avoid unnecessary computations is to keep track of the minimum cost to reach a particular node with a given number of stops. If you've already reached a node with a certain cost and a certain number of stops before, then you don't need to explore that state again.

Here's how you can optimize your code:

Add a cost array to keep track of the lowest price to reach each node for a given number of stops.
Before adding a state to the queue, check if the price is cheaper than a previously found price to get to that node with the same or fewer stops.

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        graph = collections.defaultdict(list)
        for u, v, w in flights:
            graph[u].append((v, w))

        # Initialize cost with high values
        # cost[i][j] represents the minimum cost to reach node i in j stops
        cost = [[float('inf')] * (k + 2) for _ in range(n)]
        cost[src][0] = 0

        # Queue: [(price, node, stops)]
        q = [(0, src, 0)]

        while q:
            price, node, stops = heapq.heappop(q)
            if node == dst:
                return price
            if stops > k:
                continue
            for v, w in graph[node]:
                if cost[v][stops + 1] > price + w:  # Check if we found a better price for v with stops + 1
                    cost[v][stops + 1] = price + w
                    heapq.heappush(q, (price + w, v, stops + 1))

        # We return the minimum cost among all possible stops for dst
        return min(cost[dst]) if min(cost[dst]) != float('inf') else -1

The above optimization ensures that we're not revisiting the same node-stop combination multiple times, which can reduce the number of operations and prevent a time limit exceeded error.

Bellman-Ford Algorithm

The drawing below illustrates the implementation of BFS for one layer. However, we need to execute it k + 1 times because in this context, 'k' represents a node while a layer represents an edge.

Time complexity: $O(E\cdot{K})$

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        prices = [float("inf")] * n
        prices[src] = 0
        
        for i in range(k + 1):
            tmp_prices = prices.copy()
            
            for s, d, p in flights: # source, destination, price
                if prices[s] == float("inf"):
                    continue
                if prices[s] + p < tmp_prices[d]:
                    tmp_prices[d] = prices[s] + p
            prices = tmp_prices
            
        return -1 if prices[dst] == float("inf") else prices[dst]

Reference: https://www.youtube.com/watch?v=5eIK3zUdYmE

Last updated 2 years ago