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.

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.