go: add «2699. Modify Graph Edge Weights»

Signed-off-by: Matej Focko <me@mfocko.xyz>

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{}
+}