aboutsummaryrefslogtreecommitdiff
path: root/modify_graph_edge_weight.py
blob: 67ab7c82be58bb7913ea0d6dde1210c3abf6b9cf (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
# 2699. Modify Graph Edge Weights
from collections import defaultdict
import heapq

"""
you are given an undirected weighted connected graph containing 'n' nodes
labeled from '0' to 'n - 1' and an integer array 'edges' where 'edges[i] =
[ai, bi, wi]' indicates that there is an edge between nodes 'ai' and 'bi'
with weight 'wi'. some edges have a weight of -1 ('wi = -1'), while other
have a positive weight ('wi > 0'). your task is to modify all edges with a
weight of -1 by assigning them positive integer values in the range '[1, 2 *
10^9]' so that the shortest distance between the nodes 'source' and
'destination' becomes equal to an integer 'target'. if there are multiple
modifications that make the shortest distance equal to 'target', any of them
are considered correct. return an array containing all edges in any order if
it is possible to make the shortest to equal 'target', or an empty array if
it is not possible
"""


class Solution(object):
    def modifiedGraphEdges(self, n, edges, source, destination, target):
        """
        :type n: int
        :type edges: List[List[int]]
        :type source: int
        :type destination: int
        :type target: int
        :rtype: List[List[int]]
        """
        adj = [[] for _ in range(n)]
        for u, v, w in edges:
            adj[u].append([v, w])
            adj[v].append([u, w])

        def dijkstra(source, adj, skip_neg):
            pq = [[0, source]]
            dist = defaultdict(lambda: "inf")
            dist[source] = 0
            parent = {}
            while pq:
                d, node = heapq.heappop(pq)
                if d > dist[node]:
                    continue
                for nei, w in adj[node]:
                    if w == -1:
                        if skip_neg:
                            continue
                        w = 1
                    d2 = d + w
                    if d2 < dist[nei]:
                        dist[nei] = d2
                        parent[nei] = node
                        heapq.heappush(pq, [d2, nei])
            return dist, parent

        dist_r, parent_r = dijkstra(destination, adj, skip_neg=True)
        if dist_r.get(source, "inf") < target:
            return []
        dist, parent = dijkstra(source, adj, skip_neg=False)
        if dist[destination] > target:
            return []
        path = [destination]
        while path[-1] != source:
            path.append(parent[path[-1]])
        path = path[::-1]
        edges = {(min(u, v), max(u, v)): w for u, v, w in edges}
        walked = 0
        for u, v in zip(path, path[1:]):
            e = (min(u, v), max(u, v))
            if edges[e] == -1:
                edges[e] = max(target - dist_r.get(v, "inf") - walked, 1)
                if edges[e] > 1:
                    break
            walked += edges[e]
        for e, w in edges.items():
            if w == -1:
                edges[e] = 2 * (10**9)
        return [[u, v, w] for (u, v), w in edges.items()]


if __name__ == "__main__":
    obj = Solution()
    print(
        obj.modifiedGraphEdges(
            n=5,
            edges=[[4, 1, -1], [2, 0, -1], [0, 3, -1], [4, 3, -1]],
            source=0,
            destination=1,
            target=5,
        )
    )
    print(
        obj.modifiedGraphEdges(
            n=3, edges=[[0, 1, -1], [0, 2, 5]], source=0, destination=2, target=6
        )
    )
    print(
        obj.modifiedGraphEdges(
            n=4,
            edges=[[1, 0, 4], [1, 2, 3], [2, 3, 5], [0, 3, -1]],
            source=0,
            destination=2,
            target=6,
        )
    )