细分图中的可到达节点 Dijkstra算法Python实现

题目大意

无向图（原始图）中有 n 个节点 \(u_i\)，编号从 0 到 n - 1 。对每条边 [\(u_i,v_i\)] 增加若干节点构建“细分图”，求解从节点0出发能抵达的不超过距离为maxMove的节点数量。

输入：edges = [[0,1,10],[0,2,1],[1,2,2]], maxMoves = 6, n = 3
输出：13

解题思路

求解原始图上从节点0出发能抵达的不超过距离为maxMove的节点数量，可以直接用BFS/DFS快速解决。本题增加了“细分图”的概念，在每一条边 [\(u_i,v_i\)] 上新增 \(cnt_i\) 个节点，仍然求解新的细分图上可到达节点数量。

• 分析数据规模，新的图节点数量可达到 10^4 * 10^4 = 10^8，无法直接对细分图建图，然后BFS/DFS。
• 转换一下思路，可以对原始图上的边作特殊处理，即看作路径长度，那么本题可以转化为最短路问题。
• 求解从节点0出发到其他所有节点的最短路，记录每条边两个方向走过的距离来处理额外边上的节点，统计所有总距离在最大步长范围内的点即为答案。

代码

class Solution:
    def reachableNodes(self, edges: List[List[int]], maxMoves: int, n: int) -> int:
        # 邻接矩阵
        neighbor = [[] for _ in range(n)]
        # 记录每条边的总长，[u->]方向已经访问过节点数
        new_edges = dict()

        for u, v, w in edges:
            neighbor[u].append((v, w))
            neighbor[v].append((u, w))

            new_edges[(u,v)] = [w, 0]
            new_edges[(v,u)] = [w, 0]


        from queue import PriorityQueue
        q = PriorityQueue()
        # (step, start) 优先队列，按step从小到大排序
        q.put((0, 0))
        vis = [False] * n

        ans = 0
        while not q.empty():
            step, u = q.get()

            if vis[u]:
                continue
            vis[u] = True
            ans += 1

            if step>=maxMoves:
                continue

            for v, w in neighbor[u]:
                if vis[v]:
                    # 检查边上是否有未访问的节点
                    # 剩下未走过的节点数: w-反向走过
                    left_nodes = min(w - new_edges[(v,u)][1], w-new_edges[(u,v)][1])
                    s = min(left_nodes, maxMoves-step)
                    new_edges[(u,v)][1] = s

                    ans += s
                    continue

                if step+w+1<=maxMoves:
                    ans += w
                    # print(u,v, ":", w, ans, step+w+1)
                    q.put((step+w+1, v))
                    new_edges[(u,v)][1] = w
                else:
                    # print(u,v, ":", ans, step, maxMoves-step)
                    ans += maxMoves - step
                    new_edges[(u,v)][1] = maxMoves - step
        
        return ans

后记

简简单单的想法，WA了两发才过。先是对边上节点处理不全面，需要新增走过的距离；然后是u->v的距离应为 step+w+1，put放入节点和步长位置也写反了 。

---

（完）