本质还是图的遍历

修缮管道 训练营场地的排水管道需要修缮维护。根据设计图纸可知，入水口和出水口确定了一条线段，线段上均匀坐落着N个位置固定的过滤设备，这些过滤设备沿着这条线段按照与入水口距离递增的顺序分别用1....n 来表示(1<=N<=10^3 )为了避免影响电然或其他设施，本次修缮涉及的施工均不允许超出这条线段所在的位置(无需考虑管道和过滤设备的宽度)。由于排水系统年代久远，部新旧型号(甚至同一型号)的过滤设备之间可能会不支持直接首尾相连，在入水口、出水口和所有过滤设备位置不能改变的情况下，为了尽量减小管道和后续维护的成本，我们需要根据过滤设备的兼容情况设计一人管道方案，依据该方案选择一些过滤设备并用管道按一定顺序重新连接这些过滤设备，实现水流从入水口位置的第1个过滤设备流入，经出水口位置的第N个过滤设备流出且管道总长度最短。目前过设备共有K个不同的型号(1<=K<=50)，第i个过滤设备与第j个过设备之间的距离为|i-j|个单位长度，不同型号的过滤备之间的连接兼容情况，用 KxK 的矩阵S来表示，S[i,j]= 1 代表允许第i个过设备出口通过管道连接至第j个过设备的入口，如果S[i,j] =0则代表不允许。

输入描述

第一行包含整数N和K。第二行包含N个空格分隔的整数m1...mn，分别代表各过滤设备的型号 接下来的K行代表矩阵S，每行为长度为K的字符串，S[i,j] 为第i个字符串的第j位

输出描述

请输出修缮方案所需要的管道长度的最小值，如果没有可行的方案，则输出-1

示例1

输入

6 5 1 4 2 5 3 4 10100 00010 01100 01000 11111

输出

9

说明

最优路径为1->5->3->6，管道长度为|1-5|+|5-3|+|3-6|=9

思路与代码

迪杰斯特拉求最短路。

import heapq
from collections import defaultdict
from math import inf

N, K = map(int, input().split(" "))
M = [int(c)-1 for c in input().split(" ")]

FilterIndex = defaultdict(list)
for i in range(N): FilterIndex[M[i]].append(i)

S = []
for _ in range(K):
    S.append([c for c in input()])

def dijstra():
    dis = [inf for _ in range(N)]
    dis[0] = 0
    q = [(0, 0)]  # 距离0，当前编号0
    while q:
        d, node = heapq.heappop(q)
        if node == N - 1:
            return d
        if d > dis[node]: continue

        for j in range(K):
            if S[M[node]][j] == '1':
                #当前过滤设备与 nx过滤设备是兼容的，还需要枚举当前所有的该型号的过滤设备
                for nx in FilterIndex[j]:
                    curDis = d + abs(nx - node)
                    if curDis < dis[nx]:
                        dis[nx] = curDis
                        heapq.heappush(q, (curDis, nx))

    return -1

print(dijstra())