[leetcode 2846][LCA]边权重均等查询

题目链接

leetcode 边权重均等查询

前置概念

LCA倍增树算法

链接

解题思路

根据LCA的的性质，要确定两个节点的最短路径只需要确定两个节点与根节点的最短路径以及两个节点最近公共祖先与根节点的最短路径即可，所以只需要初始化一个树，再预处理计算出每个节点到根节点经过的边权为x的边的数量就能得出答案

先将节点进行预处理，随机选择一个节点作为根节点，再计算出每个节点的父节点和子节点 

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
   //e[i]代表与节点i连接的节点以及边权(状态压缩)
var e = make([][]int, n)
for _, v := range edges {
	i, j, w := v[0], v[1], v[2]
	e[i] = append(e[i], j*27+w)
	e[j] = append(e[j], i*27+w)
}
   //父节点，parent[i]表示i的父节点以及边权
var parent = make([][]int, n)
var child = make([][]int, n)//子节点
   var isVis = make([]bool, n)
   var dep = make([]int, n)//节点深度
var treeInit func(int, int)
   //从根节点开始遍历
treeInit = func(i int, d int) {
	next := e[i]
	dep[i] = d
	for _, v := range next {
		j, w := v/27, v%27
           //若节点j与i相连且j未被访问则j为i的子节点
		if !isVis[j] {
			isVis[j] = true
			parent[j] = []int{i, w}
			child[i] = append(child[i], j)
			treeInit(j, d+1)
		}
	}
}
   isVis[n / 2] = true //根节点可以为任意节点
   treeInit(n / 2, 0)
初始化树之后初始化LCA所需的pa数组，并声明lca函数 
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
   var pa = make([][]int, n)
var branch = make([]int, n)
var initPa func(int)
initPa = func(i int) {
	d := dep[i]
	branch[d] = i
	for x := 0; d-(1<<x) >= 0; x++ {
		pa[i] = append(pa[i], branch[d-(1<<x)])
	}
	for _, v := range child[i] {
		InitPa(v)
	}
}
InitPa(n / 2)
   var lca = func(i, j int) int {
	if dep[i] > dep[j] {
		h := dep[i] - dep[j]
		m := bits.Len(uint(h))
		for k := range m {
			if (h >> k & 1) == 1 {
				i = pa[i][k]
			}
		}
	} else {
		h := dep[j] - dep[i]
		m := bits.Len(uint(h))
		for k := range m {
			if (h >> k & 1) == 1 {
				j = pa[j][k]
			}
		}
	}
	if i == j {
		return i
	}
	for x := len(pa[i]) - 1; x >= 0; x-- {
		if pa[i][x] != pa[j][x] {
			i, j = pa[i][x], pa[j][x]
               x = len(pa[i])
		}
	}
	return pa[i][0]
}
声明一个二维数组weight，weight[i][n]表示节点i到根节点经过的边权为n的边的数量 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
var weight = make([][]int, n)
for i := range weight {
       //边权最大为26
	weight[i] = make([]int, 27)
}
var weightInit func(int)
weightInit = func(i int) {
	if parent[i] != nil {
		p, w := parent[i][0], parent[i][1]
		copy(weight[i], weight[p])
		weight[i][w]++
	}
	for _, v := range child[i] {
		weightInit(v)
	}
}
weightInit(n / 2)
最后遍历queries 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
   var ret []int
for _, q := range queries {
	i, j := q[0], q[1]
	wi, wj := weight[i], weight[j]
	wp := weight[lca(i, j)] //找出i和j的最近公共节点
	sum := 0
	maxS := 0
       //d(i, j) = h(i) + h(j) - 2 * h(lca(i, j))
	for k := 0; k <= 26; k++ {
		s := wi[k] + wj[k] - 2*wp[k]
		sum += s
		maxS = max(maxS, s)
	}
	ret = append(ret, sum-maxS)
}
return ret

上述代码，树的构建时间复杂度为，pa数组的初始化时间复杂度为，weight数组初始化的时间复杂度为，每次查询lca的复杂度为，总查询时间复杂度为，所以整个代码的时间复杂度为