百度9-3号笔试代码题

grid = [[1,2,2],
        [4,5,2],
        [7,8,2]]

n = len(grid)
table = {}

for i in range(n):
    for j in range(n):
        if (i,j) not in table:
            table[(i,j)] = []

        if 0 <= i+1 < n:
            table[(i,j)].append([(i+1,j),abs(grid[i+1][j] - grid[i][j])])
        if 0 <= i-1 < n:
            table[(i,j)].append([(i-1,j),abs(grid[i-1][j] - grid[i][j])])
        if 0 <= j+1 < n:
            table[(i,j)].append([(i,j+1),abs(grid[i][j+1] - grid[i][j])])
        if 0 <= j-1 < n:
            table[(i,j)].append([(i,j-1),abs(grid[i][j-1] - grid[i][j])])

visited = [[0 for _ in range(n)] for _ in range(n)]

dist = [[999 for _ in range(n)] for _ in range(n)]

dist[0][0] = 0

for _ in range(n*n):
    my_min = 1000
    pivot = (0,0)

    for i in range(n):
        for j in range(n):
            if visited[i][j] == 0:
                if dist[i][j] < my_min:
                    my_min = dist[i][j]
                    pivot = (i,j)

    idx,idy = pivot
    visited[idx][idy] = 1
    path = table[pivot]

    for m in range(len(path)):
        next_ptx,next_pty = path[m][0]
        next_dis = path[m][1]
        dist[next_ptx][next_pty] = min(dist[next_ptx][next_pty],dist[idx][idy] + next_dis)

print(dist)