LeetCode/go/minimum-cost-walk-in-weighted-graph.go

package main

type Edge struct {
	v      int
	weight int
}

type Component struct {
	id   int
	cost int
}

func minimumCost(n int, edges, query [][]int) []int {
	// reconstruct graph
	makeGraph := func() [][]Edge {
		g := make([][]Edge, n)
		for _, edge := range edges {
			u, v, weight := edge[0], edge[1], edge[2]

			g[u] = append(g[u], Edge{v: v, weight: weight})
			g[v] = append(g[v], Edge{v: u, weight: weight})
		}

		return g
	}
	g := makeGraph()

	visited := make([]bool, n)
	components := make([]int, n)
	costs := make([]int, 0)

	var getCost func(int, int) int
	getCost = func(id, u int) int {
		cost := int((^uint(0)) >> 1)

		components[u] = id
		visited[u] = true

		for _, edge := range g[u] {
			cost &= edge.weight

			if visited[edge.v] {
				continue
			}

			cost &= getCost(id, edge.v)
		}

		return cost
	}

	// traverse components
	componentId := 0
	for u := range n {
		if visited[u] {
			continue
		}

		costs = append(costs, getCost(componentId, u))
		componentId++
	}

	// construct answer to queries
	answer := make([]int, len(query))
	for i, q := range query {
		start, end := q[0], q[1]

		if components[start] == components[end] {
			answer[i] = costs[components[start]]
		} else {
			answer[i] = -1
		}
	}

	return answer
}