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:
(neighbor, edge_cost) pairs for the given node.current_dist + edge_cost. For problems like 1631, you might be taking the max instead of sum.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.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
LeetCode 743: Network Delay Time
current_dist + edge_costdef 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
LeetCode 1514: Path with Maximum Probability
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])
LeetCode 1631: Path With Minimum Effort
max(current_effort, abs(height_diff)).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.