LC 2699. 修改图中的边权

https://leetcode.cn/problems/modify-graph-edge-weights/

完全没有啥思路,看了题目的几个提示之后想到了一个解法,也是这个 题解 里面提到的错误解法。

大致思路就是要做两边最短路:

第一次找destionation -> x的最短路,这里面记录为D1

第二次找source -> x的最短路,并且期间需要使用D1的结果

但是上面的修改并不一定保证,最后的最短路是 src -> x -> y -> dest. 因为在第一次最短路里面dest->src并不一定经过(x,y)这个节点。

我觉得可能这题最关键就是这里,题解里面提到的:

如果 y 不在最短路上,修改 x - y 并不会对最短路产生影响,所以代码中并没有判断 y 是否在最短路上。

我们每次发现这些边的时候直接去修改就好了,不用判断是否在最短路上。我们假设在最短路上,修改完成之后做一次检查,判断是否满足target就好。

这种两次最短路的算法直接没有看过,比较新鲜。我的代码里面不好直接测试样例,所以增加了一个 `DOCHECK` 标记。

class Solution:
    def modifiedGraphEdges(self, n: int, edges: List[List[int]], source: int, destination: int, target: int) -> List[
        List[int]]:

        # build adj on fixed cost
        # assign varied cost with 1.
        adj = [[] for _ in range(n)]
        disconnect = set()
        for a, b, w in edges:
            if w == -1:
                disconnect.add((a, b))
                disconnect.add((b, a))
                w = 1
            adj[a].append((b, w))
            adj[b].append((a, w))

        def shortest(s, e, adj):
            visit = [-1] * n
            import heapq
            hp = []
            hp.append((0, s, -1))
            while hp:
                (d, x, p) = heapq.heappop(hp)
                if visit[x] != -1: continue
                visit[x] = d
                if x == e: break
                for y, w in adj[x]:
                    if visit[y] != -1: continue
                    heapq.heappush(hp, (d + w, y, x))
            return visit

        if True:
            D1 = shortest(destination, source, adj)
            d = D1[source]
            if d > target:
                return []

        # dijkstra again with deterministic order.
        def fixedWeight(s, e, adj, D):
            visit = [-1] * n
            hp = []
            hp.append((0, s, -1, 0))
            weight = {}
            while hp:
                (d, x, p, w) = heapq.heappop(hp)
                if visit[x] != -1: continue
                visit[x] = d
                if p != -1:
                    weight[(x, p)] = w
                if x == e: break
                for y, w in adj[x]:
                    if (x, y) in disconnect:
                        if D[y] != -1 and (d + D[y] + w) <= target:
                            w = target - d - D[y]
                        else:
                            w = target + 1
                    if visit[y] != -1: continue
                    heapq.heappush(hp, (d + w, y, x, w))
            return visit, weight

        D2, W = fixedWeight(source, destination, adj, D1)
        if D2[destination] != target:
            return []

        output = []
        for a, b, w in edges:
            if w != -1:
                output.append((a, b, w))
                continue

            w = W.get((a, b)) or W.get((b, a)) or w
            if w == -1:
                w = target + 1
            output.append((a, b, w))
        output = [list(x) for x in output]

        DOCHECK = False
        if os.environ.get('USER') == 'dirlt':
            # print('mother')
            DOCHECK = True

        if DOCHECK:
            def check(edges):
                adj = [[] for _ in range(n)]
                for a, b, w in edges:
                    adj[a].append((b, w))
                    adj[b].append((a, w))
                D = shortest(source, destination, adj)
                if D[destination] != target:
                    print('FAILED', D[destination], target)

            check(output)

        return output