LeetCode/go/path-with-maximum-probability.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
}