JZX轻语：简
LeetCode 2959 - 访问消失节点的最少时间
发表于2024年07月17日
多源最短路径问题的变种，可使用Floyd算法解决。题意要求求解的是去掉某些节点后，所有节点之间的最短路径都不超过maxDistance的方案数。可以逆向思考，求解子集中所有节点之间的最短路径都不超过maxDistance的子集数量。因此，我们可以使用回溯法枚举所有子集，然后使用Floyd算法求解子集中所有节点之间的最短路径，最后统计满足条件的子集数量。需要注意的是，如果在每个枚举中都重新来一遍Floyd算法，时间开销会特别巨大。
我们可以回顾下Floyd算法的原理，本质是动态规划，在每次迭代中，我们只需要关注新增的节点(允许经过该节点)对所有节点间的最短路径的影响。对于本题目，我们可以在参数中传递前面的节点集合中的最短路径矩阵，然后在本次枚举中，只需要关注新增的节点对前面节点间的最短路径的影响即可。
import copy
class Solution:
def numberOfSets(self, n: int, maxDistance: int, roads: List[List[int]]) -> int:
ans = 1
# 初始化邻接矩阵, 本质其实就是Floyd算法的初始化步骤
matrix = [[-1] * n for _ in range(n)]
for u, v, w in roads:
if matrix[u][v] == -1 or w < matrix[u][v]:
matrix[u][v] = matrix[v][u] = w
for i in range(n):
matrix[i][i] = 0
dis = matrix
def traceback(cur_set: set[int], k: int, old_dis: List[List[int]]):
nonlocal ans
# 新的最短路径矩阵, 需要深拷贝, 否则会影响到上一次的结果
new_dis = copy.deepcopy(old_dis)
# 记录当前枚举中的最大距离
max_dis = -1
# 另外, 还需要判断是否子集内所有节点都连通
all_connected = True
# 类似于Floyd算法的内层迭代过程
# 需要注意的是，所有的节点距离都要更新
# (后面的枚举会用到这些中间状态的!比如后面的枚举用到了节点l，当此时节点l不在集合内，而l到所有节点的距离都满足条件，如果不更新则不知道上述事实),
# 不能只更新集合内节点间的距离
for i in range(n):
for j in range(i, n):
# 如果i-k和k-j都是连通的, 且i-k-j的距离小于i-j的距离, 则更新i-j的距离
if old_dis[i][k] != -1 and old_dis[k][j] != -1 and (
old_dis[i][j] == -1 or old_dis[i][j] > old_dis[i][k] + old_dis[k][j]
):
# 需要注意i->j和j->i的距离都要更新
new_dis[i][j] = new_dis[j][i] = old_dis[i][k] + old_dis[k][j]
if i in cur_set and j in cur_set:
# 只记录集合内节点间的最大距离，以及它们是否连通
if new_dis[i][j] == -1:
all_connected = False
else:
max_dis = max(max_dis, new_dis[i][j])
if max_dis != -1 and max_dis <= maxDistance and all_connected:
ans += 1
# 需要注意的是，如果当前枚举的节点集合不满足条件，后续也需要枚举的，因此可能后面引入的节点会使得一些距离变得更短
for new_k in range(k + 1, n):
cur_set.add(new_k)
traceback(cur_set, new_k, new_dis)
cur_set.remove(new_k)
for i in range(n):
traceback({i}, i, dis)
return ans