LeetCode/go/modify-graph-edge-weights.go

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