60 lines
1.1 KiB
Go
60 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
|
||
|
}
|