Ming
0%

LeetCode-399

Posted on Edited on In

题目

399. 除法求值

结果

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
func calcEquation(equations [][]string, values []float64, queries [][]string) (ans []float64) {
	// create UnionFind
	u := UnionFind{make(map[string]string), make(map[string]float64)}
	for i, v := range equations {
		// new node initialization
		if _, ok := u.parent[v[0]]; !ok {
			u.parent[v[0]] = v[0]
			u.weight[v[0]] = 1
		}
		if _, ok := u.parent[v[1]]; !ok {
			u.parent[v[1]] = v[1]
			u.weight[v[1]] = 1
		}
		u.union(v[0], v[1], values[i])
	}

	for _, v := range queries {
		_, ok1 := u.parent[v[0]]
		_, ok2 := u.parent[v[1]]
		if !ok1 || !ok2 {
			ans = append(ans, -1)
		} else {
			ans = append(ans, u.ratio(v[0], v[1]))
		}
	}
	return
}

type UnionFind struct {
	parent map[string]string // whosyourdaddy
	weight map[string]float64
}

func (u *UnionFind) union(x, y string, val float64) {
	rootX, rootY := u.find(x), u.find(y)
	if rootX == rootY {
		return
	}
	u.parent[rootX] = rootY
	u.weight[rootX] = u.weight[y] * val / u.weight[x]
}

// find id of the root node
func (u *UnionFind) find(x string) string {
	// checkout if the current node is the root node
	if x != u.parent[x] {
		// find root node recursively
		origin := u.parent[x]
		u.parent[x] = u.find(u.parent[x])
		u.weight[x] *= u.weight[origin]
	}
	return u.parent[x]
}

func (u *UnionFind) ratio(x, y string) float64 {
	// if two nodes are in the same tree set
	rootX, rootY := u.find(x), u.find(y)
	if rootX == rootY {
		// because tree height is 2
		return u.weight[x] / u.weight[y]
	}
	return -1
}