[Pattern&Template] Dijkstra Algorithm (applicable to question 743, 1514 and 1631)

Below is a single, flexible Dijkstra’s algorithm implementation that can be adapted for different problem setups like LeetCode 743, 1514, and 1631. The main idea is to write a generic dijkstra function that takes in a get_neighbors function and a combine_cost function, so you can customize the graph behavior and cost calculation logic for different problems.

How This Works:

  • get_neighbors(node): A function you provide that returns a list of (neighbor, edge_cost) pairs for the given node.
  • combine_cost(current_dist, edge_cost): A function you provide that calculates the new distance after taking an edge from the current node. For standard shortest paths, it might be current_dist + edge_cost. For problems like 1631, you might be taking the max instead of sum.
  • The dijkstra function itself is generic and doesn’t assume a specific cost structure. It simply uses a min-heap to pick the next best candidate and relax edges according to the given logic.

Unified Dijkstra Template

import heapq

def dijkstra(start, n, get_neighbors, combine_cost, initial_dist=float('inf')):
    """
    :param start: starting node (or cell)
    :param n: number of nodes (or size for indexing dist if needed)
    :param get_neighbors: function(node) -> List[(neighbor, edge_cost)]
    :param combine_cost: function(current_dist, edge_cost) -> new_dist (defines how to update distances)
    :param initial_dist: initial distance for all nodes (default inf)
    :return: dist array with shortest "distances" as defined by combine_cost
    """
    dist = [initial_dist]*(n)
    dist[start] = 0
    pq = [(0, start)]  # (distance, node)

    while pq:
        current_dist, node = heapq.heappop(pq)
        if current_dist > dist[node]:
            continue
        for (nei, edge_cost) in get_neighbors(node):
            new_dist = combine_cost(dist[node], edge_cost)
            if new_dist < dist[nei]:
                dist[nei] = new_dist
                heapq.heappush(pq, (new_dist, nei))
    return dist

Applying to Each Problem

  1. LeetCode 743: Network Delay Time

    • Graph: Weighted directed edges with non-negative weights.
    • Combine Cost: current_dist + edge_cost
    • get_neighbors: Provided by adjacency list.
    def networkDelayTime(times, n, k):
        # build adjacency list
        adj = [[] for _ in range(n)]
        for u,v,w in times:
            adj[u-1].append((v-1, w))
    
        def get_neighbors(node):
            return adj[node]
    
        def combine_cost(current_dist, edge_cost):
            return current_dist + edge_cost
    
        dist = dijkstra(start=k-1, n=n, get_neighbors=get_neighbors, combine_cost=combine_cost)
        ans = max(dist)
        return ans if ans < float('inf') else -1
    
  2. LeetCode 1514: Path with Maximum Probability

    • Graph: Edges have probabilities; we want to maximize probability.
    • Trick: Instead of maximizing product of probabilities, we minimize negative log or equivalently just store probabilities and use a max-heap approach. But to unify with min-heap Dijkstra, we can store log(prob) as the cost, or simply store negative probabilities to convert max problem into a min problem.Method: Use log(prob) so that a larger probability results in a smaller value.
    import math
    
    def maxProbability(n, edges, succProb, start, end):
        adj = [[] for _ in range(n)]
        for (u,v), p in zip(edges, succProb):
            # cost = -log(p)
            cost = -math.log(p) if p > 0 else float('inf')
            adj[u].append((v, cost))
            adj[v].append((u, cost))
    
        def get_neighbors(node):
            return adj[node]
    
        def combine_cost(current_dist, edge_cost):
            return current_dist + edge_cost
    
        dist = dijkstra(start, n, get_neighbors, combine_cost)
    
        # dist[end] is the sum of -log(probabilities)
        # probability = exp(-dist[end])
        return 0.0 if dist[end] == float('inf') else math.exp(-dist[end])
    
  3. LeetCode 1631: Path With Minimum Effort

    • We have a grid. Each edge cost is defined by max(current_effort, abs(height_diff)).
    • This means we don’t sum weights; we take the maximum of the current distance and the new edge cost.
    def minimumEffortPath(heights):
        rows, cols = len(heights), len(heights[0])
        def in_bounds(r, c):
            return 0 <= r < rows and 0 <= c < cols
    
        directions = [(1,0),(-1,0),(0,1),(0,-1)]
    
        def get_neighbors(cell):
            r, c = divmod(cell, cols)
            result = []
            for dr, dc in directions:
                nr, nc = r+dr, c+dc
                if in_bounds(nr, nc):
                    # edge cost = abs difference
                    edge_cost = abs(heights[nr][nc] - heights[r][c])
                    nei = nr*cols + nc
                    result.append((nei, edge_cost))
            return result
    
        def combine_cost(current_dist, edge_cost):
            # Take the max, not sum
            return max(current_dist, edge_cost)
    
        dist = dijkstra(start=0, n=rows*cols, get_neighbors=get_neighbors, combine_cost=combine_cost)
        return dist[rows*cols - 1]
    

With this single dijkstra function, we’ve shown how each problem (743, 1514, and 1631) can be solved by plugging in the appropriate get_neighbors and combine_cost logic, along with constructing the right graph representation. The key difference in each is how we define and interpret edge costs and how we combine them with the current distance.

This unified approach highlights how Dijkstra’s algorithm can be adapted for various shortest path scenarios by customizing the cost function and neighbor retrieval.

Comments (0)