Optimizing Dijkstra on grid Python (Microsoft Interview)

Question

Given a square grid of size N, each cell of which contains integer cost which represents a cost to traverse through that cell, we need to find a path from top left cell to bottom right cell by which the total cost incurred is minimum. From the cell (i,j) we can go (i,j-1), (i, j+1), (i-1, j), (i+1, j).

Note: It is assumed that negative cost cycles do not exist in the input matrix.

Example 1:
Input: grid = {{9,4,9,9},{6,7,6,4},
{8,3,3,7},{7,4,9,10}}
Output: 43
Explanation: The grid is-
9 4 9 9
6 7 6 4
8 3 3 7
7 4 9 10
The minimum cost is-
9 +たす 4 +たす 7 +たす 3 +たす 3 +たす 7 +たす 10 =わ 43.

My working code:

global INT_MAX
# value of n is passed by the driver along with the grid
INT_MAX = 3 ** 38
class Solution:
 
 # <- this returns the minVertex's co-orinates ->
 def minimumVertex(self,distMat,visited):
 minVertexi = 0
 minVertexj = 0
 max1 = INT_MAX
 for i in range(n):
 for j in range(n):
 if distMat[i][j] < max1 and visited[i][j] == False:
 max1 = distMat[i][j]
 minVertexi = i
 minVertexj = j
 return minVertexi, minVertexj
 
 def minimumCostPath(self, grid):
 
 # <- visited matrix is created with all False ->
 visited = []
 for i in range(n):
 eachRow = []
 for j in range(n):
 eachRow.append(False)
 visited.append(eachRow)
 # <- distMat is created with all values initialized with INT_MAX -> 
 distMat = []
 for i in range(n):
 eachRow1 = []
 for j in range(n):
 eachRow1.append(INT_MAX)
 distMat.append(eachRow1)
 distMat[0][0] = grid[0][0]
 drow = [-1, 1, 0, 0]
 dcol = [0, 0, -1, 1]
 
 for _ in range(n):
 for currj in range(n):
 mini,minj = self.minimumVertex(distMat,visited)
 
 visited[mini][minj] = True
 for k in range(len(drow)):
 nbri = mini + drow[k]
 nbrj = minj + dcol[k]
 # print("outside nbri = ",nbri,"outside nbrj = ",nbrj)
 if nbri >=0 and nbrj >= 0 and nbri < n and nbrj < n:
 if visited[nbri][nbrj] == False:
 # print("inside nbri = ",nbri,"inside nbrj = ",nbrj)
 if distMat[nbri][nbrj] > distMat[mini][minj] + grid[nbri][nbrj]:
 distMat[nbri][nbrj] = distMat[mini][minj] + grid[nbri][nbrj]
 return distMat[n-1][n-1]
 

Expected Time Compelxity: O(n2*log(n))

Expected Auxiliary Space: O(n2)

Problem with my code:

My logic is alright however it is exceeding the time limit, how do I optimize it?

NOTE: The grid is always a square grid and N (dimension) is passed to the function along with the grid. I took N = 4 here because I was debugging the code for a single test case.

• RootTwo: I believe this solution is O(n^4). minimumVertex() is O(n^2) and it is called n^2 times in the nested loops in minimumCostPath(). Try using a heap to keep track of the minimum vertex.

Answer 1

Actual time complexity

As you correctly state, the expected time complexity of Dijkstra's on a square matrix is O(n^2 log n). However, your algorithm has a complexity of at least O(n^4) due to the way you locate the minimum vertex.

for _ in range(n):
 for currj in range(n):
 # <--- executed n^2 times here
 mini,minj = self.minimumVertex(distMat,visited)
 ...
def minimumVertex(self,distMat,visited):
 minVertexi = 0
 minVertexj = 0
 max1 = INT_MAX
 for i in range(n):
 for j in range(n):
 # <--- takes n^2 time
 ...

In fact, the algorithm you implemented isn't purely Dijkstra's, but a blend of Dijkstra's algorithm and Floyd's algorithm. You should read through the links above, particularly the Pseudocode sections to understand how each of the algorithms works, and the parts that you're mixing up.

Dijkstra's algorithm

The key data structure which allows Dijkstra's algorithm to work in O(n^2 log n) time is a priority queue which allows you to insert elements and remove the top element in O(log n) time. This data structure is commonly implemented using a binary heap, and has a standard implementation in Python as the heapq library.

You should adapt your code to use this data structure instead of the distance table as shown below. It is important that the first field of the tuple is the distance, since this determines the sorting.

def minimumCostPath(self, grid):
 # definitions of n, visited, drow, dcol
 queue = [(grid[0][0], 0, 0)]
 while queue:
 dist, i, j = heapq.heappop(queue)
 if visited[i][j]:
 continue
 visited[i][j] = True
 
 if i == j == n - 1:
 return dist
 for k in range(len(drow)):
 ni = i + drow[k]
 nj = j + dcol[k]
 if 0 <= ni < n and 0 <= nj < n:
 heapq.heappush(queue, (dist + grid[ni][nj], ni, nj))
 assert False, "queue is exhausted but we haven't reached the destination yet"

Benchmarking surprise

After writing all of this, I decided to quantify the improvement by benchmarking. This is the code I used to see how my implementation compares against the original.

if __name__ == "__main__":
 import time
 N = 100
 grid = [[j for _ in range(N)] for j in range(N)]
 solver = Solution()
 start = time.time()
 print(solver.minimumCostPath(grid))
 end = time.time()
 print(f"Solved grid with size {N} in {1000*(end - start):.3f}ms")

To my surprise, your version was actually faster, but more importantly it gave incorrect results.

My logic is alright

No, it's not. The issue is the usage of n in the minimumVertex function. Because you define n = 4 as a global, this value is used instead of the actual size of the grid. This makes your algorithm appear fast, but actually be incorrect for n > 4.

Conclusion

As it turns out, while the algorithm you (I imagine) intended to implement is O(n^4), your code was actually O(n^2) all along. If this really resulted in TLE, Python is not going to be able to handle this question. What I find much more likely is that your code failed due to an incorrect response - using the correct implementation of Dijkstra's above is likely to fix this.

• Shubham Prashar: This code was meant for an online judge and the reason I defined n = 4 because I was debugging the code for a single test case where N was 4. Now understand this, the grid is always a square grid (rows = columns) and the driver calls my function with the grid and N so that's really not an issue here, I hope you are getting what I am trying to say.
• Shubham Prashar: Also, no offence but refrain from saying the logic is wrong and try asking the OP or thinking the reason behind something.
• apilat: @ShubhamPrashar Next time please include the whole code in the question. In minimumCostPath you calculate the correct n from the given grid, but the proceed to use a different n in minimumVertex. With the snippet provided there is no reasonable explanation other than a logic error. In fact, even if you define the global n to be the size of the grid, I would still consider this a logic error - you should either use the global n in both functions, or the calculated n in both functions.
• Shubham Prashar: Yeah, I totally agree with you here but what do you think about the complexity part? Everythings fine yet getting the TLE. How to optimize it?
• apilat: @ShubhamPrashar Re-read the first two sections of my answer, where I show that the complexity of your solution is actually O(n^4). You should use a correct implementation of Dijkstra's, which will give you the desired complexity of O(n^2 log n).