LeetCode/go/path-with-maximum-probability.go
Matej Focko f848a7b9c7
go: add «1514. Path with Maximum Probability»
Signed-off-by: Matej Focko <me@mfocko.xyz>
2024-08-31 21:32:31 +02:00

59 lines
1.1 KiB
Go

package main
import (
"cmp"
pq "github.com/emirpasic/gods/v2/queues/priorityqueue"
)
func maxProbability(n int, edges [][]int, succProb []float64, start int, end int) float64 {
type Edge struct {
v int
p float64
}
makeGraph := func() [][]Edge {
g := make([][]Edge, n)
for i, edge := range edges {
u, v, p := edge[0], edge[1], succProb[i]
g[u] = append(g[u], Edge{v: v, p: p})
g[v] = append(g[v], Edge{v: u, p: p})
}
return g
}
g := makeGraph()
type Entry struct {
probability float64
u int
}
byProbability := func(a, b Entry) int {
return -cmp.Compare(a.probability, b.probability)
}
maxP := make([]float64, n)
q := pq.NewWith(byProbability)
maxP[start] = 1
q.Enqueue(Entry{probability: maxP[start], u: start})
for candidate, ok := q.Dequeue(); ok; candidate, ok = q.Dequeue() {
if candidate.u == end {
return candidate.probability
}
for _, edge := range g[candidate.u] {
p := candidate.probability * edge.p
if p > maxP[edge.v] {
maxP[edge.v] = p
q.Enqueue(Entry{probability: p, u: edge.v})
}
}
}
return 0
}