diff --git a/go/modify-graph-edge-weights.go b/go/modify-graph-edge-weights.go new file mode 100644 index 0000000..11600ae --- /dev/null +++ b/go/modify-graph-edge-weights.go @@ -0,0 +1,108 @@ +package main + +import ( + "cmp" + + pq "github.com/emirpasic/gods/v2/queues/priorityqueue" +) + +const ( + INFINITY = 2_000_000_000 +) + +func modifiedGraphEdges(n int, edges [][]int, source int, destination int, target int) [][]int { + type Edge struct { + v int + weight int + } + + makeGraph := func() [][]Edge { + graph := make([][]Edge, n) + + for _, edge := range edges { + u, v, weight := edge[0], edge[1], edge[2] + if weight == -1 { + continue + } + + graph[u] = append(graph[u], Edge{v: v, weight: weight}) + graph[v] = append(graph[v], Edge{v: u, weight: weight}) + } + + return graph + } + g := makeGraph() + + type DijkstraEntry struct { + distance int + u int + } + byDistance := func(a, b DijkstraEntry) int { + return cmp.Compare(a.distance, b.distance) + } + + dijkstra := func() int { + minDistance := make([]int, n) + for i, _ := range minDistance { + minDistance[i] = INFINITY + } + + q := pq.NewWith[DijkstraEntry](byDistance) + + minDistance[source] = 0 + q.Enqueue(DijkstraEntry{distance: 0, u: source}) + + for next, ok := q.Dequeue(); ok; next, ok = q.Dequeue() { + if next.distance > minDistance[next.u] { + continue + } + + for _, edge := range g[next.u] { + if next.distance+edge.weight < minDistance[edge.v] { + minDistance[edge.v] = next.distance + edge.weight + q.Enqueue(DijkstraEntry{distance: minDistance[edge.v], u: edge.v}) + } + } + } + + return minDistance[destination] + } + + distance := dijkstra() + if distance < target { + return [][]int{} + } + + matches := (distance == target) + + for _, edge := range edges { + u, v, weight := edge[0], edge[1], edge[2] + if weight != -1 { + continue + } + + if matches { + weight = INFINITY + } else { + weight = 1 + } + g[u] = append(g[u], Edge{v: v, weight: weight}) + g[v] = append(g[v], Edge{v: u, weight: weight}) + + if !matches { + newDistance := dijkstra() + + if newDistance <= target { + matches = true + weight += target - newDistance + } + } + + edge[2] = weight + } + + if matches { + return edges + } + return [][]int{} +}