109 lines
2 KiB
Go
109 lines
2 KiB
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{}
|
||
|
|
}
|