LeetCode/go/number-of-ways-to-arrive-at-destination.go

package main

func countPaths(n int, roads [][]int) int {
	type Entry struct {
		time uint64
		ways uint64
	}

	const MOD uint64 = 1_000_000_007

	dp := make([][]Entry, n)
	for i := range n {
		dp[i] = make([]Entry, n)
	}

	// Initialize for trivial values
	for p := 0; p < n*n; p++ {
		u, v := p/n, p%n

		if u != v {
			dp[u][v].time = uint64(1e12)
			dp[u][v].ways = 0
		} else {
			dp[u][v].time = 0
			dp[u][v].ways = 1
		}
	}

	// add edges
	for _, road := range roads {
		u, v, t := road[0], road[1], road[2]

		// time
		dp[u][v].time = uint64(t)
		dp[v][u].time = uint64(t)

		// ways
		dp[u][v].ways = 1
		dp[v][u].ways = 1
	}

	// Floyd-Warshall
	for t := 0; t < n*n*n; t++ {
		v, u, w := (t/n)/n, (t/n)%n, t%n

		if u != v && v != w {
			newTime := dp[u][v].time + dp[v][w].time

			if newTime < dp[u][w].time {
				// shorter path
				dp[u][w].time = newTime
				dp[u][w].ways = (dp[u][v].ways * dp[v][w].ways) % MOD
			} else if newTime == dp[u][w].time {
				// same time, but another way
				dp[u][w].ways = (dp[u][w].ways + dp[u][v].ways*dp[v][w].ways) % MOD
			}
		}
	}

	return int(dp[n-1][0].ways)
}