2021-07-31:给定数组father,大小为N,表示一共有N个节点,father[i] = j 表示点i的父亲是点j, father表示的树一定是一棵树而不是森林,给定数组values,大小为N,values[i]=v表示节点i的权值是v。实现如下4个方法,保证4个方法都很快!1)让某个子树所有节点值加上v,入参:int head, int v;2)查询某个子树所有节点值的累加和,入参:int head;3)在树上从a到b的整条链上所有加上v,入参:int a, int b, int v;4)查询在树上从a到b的整条链上所有节点值的累加和,入参:int a, int b。
福大大 答案2021-07-31:
时间紧。见代码。
代码用golang编写。代码如下:
package main import ( "fmt" "math" ) func main() { father := []int{1, 2, 3, 4, 5} values := []int{6, 7, 8, 9, 10} tc := NewTreeChain(father, values) right := NewRight(father, values) tc.addSubtree(11, 22) right.addSubtree(11, 22) tc.addSubtree(33, 44) right.addSubtree(33, 44) if tc.queryChain(33, 44) != right.queryChain(33, 44) { fmt.Println("错误") } else { fmt.Println("正确") } } type TreeChain struct { // 时间戳 0 1 2 3 4 tim int // 节点个数是n,节点编号是1~n n int // 谁是头 h int // 朴素树结构 tree [][]int // 权重数组 原始的0节点权重是6 -> val[1] = 6 val []int // father数组一个平移,因为标号要+1 fa []int // 深度数组! dep []int // son[i] = 0 i这个节点,没有儿子 // son[i] != 0 j i这个节点,重儿子是j son []int // siz[i] i这个节点为头的子树,有多少个节点 siz []int // top[i] = j i这个节点,所在的重链,头是j top []int // dfn[i] = j i这个节点,在dfs序中是第j个 dfn []int // 如果原来的节点a,权重是10 // 如果a节点在dfs序中是第5个节点, tnw[5] = 10 tnw []int // 线段树,在tnw上,玩连续的区间查询或者更新 seg *SegmentTree } func NewTreeChain(father []int, values []int) *TreeChain { ret := &TreeChain{} // 原始的树 tree,弄好了,可以从i这个点,找到下级的直接孩子 // 上面的一大堆结构,准备好了空间,values -> val // 找到头部点 ret.initTree(father, values) // fa; // dep; // son; // siz; ret.dfs1(ret.h, 0) // top; // dfn; // tnw; ret.dfs2(ret.h, ret.h) ret.seg = NewSegmentTree(ret.tnw) ret.seg.build(1, ret.n, 1) return ret } func (this *TreeChain) initTree(father []int, values []int) { this.tim = 0 this.n = len(father) + 1 this.tree = make([][]int, this.n) this.val = make([]int, this.n) this.fa = make([]int, this.n) this.dep = make([]int, this.n) this.son = make([]int, this.n) this.siz = make([]int, this.n) this.top = make([]int, this.n) this.dfn = make([]int, this.n) this.tnw = make([]int, this.n) this.n-- cnum := make([]int, this.n) for i := 0; i < this.n; i++ { this.val[i+1] = values[i] } for i := 0; i < this.n; i++ { if father[i] == i { this.h = i + 1 } else { cnum[father[i]]++ } } this.tree[0] = make([]int, 0) for i := 0; i < this.n; i++ { this.tree[i+1] = make([]int, cnum[i]) } for i := 0; i < this.n; i++ { if i+1 != this.h { cnum[father[i]]-- this.tree[father[i]+1][cnum[father[i]]] = i + 1 } } } // u 当前节点 // f u的父节点 func (this *TreeChain) dfs1(u int, f int) { this.fa[u] = f this.dep[u] = this.dep[f] + 1 this.siz[u] = 1 maxSize := -1 for _, v := range this.tree[u] { // 遍历u节点,所有的直接孩子 this.dfs1(v, u) this.siz[u] += this.siz[v] if this.siz[v] > maxSize { maxSize = this.siz[v] this.son[u] = v } } } // u当前节点 // t是u所在重链的头部 func (this *TreeChain) dfs2(u int, t int) { this.tim++ this.dfn[u] = this.tim this.top[u] = t this.tnw[this.tim] = this.val[u] if this.son[u] != 0 { // 如果u有儿子 siz[u] > 1 this.dfs2(this.son[u], t) for _, v := range this.tree[u] { if v != this.son[u] { this.dfs2(v, v) } } } } // head为头的子树上,所有节点值+value // 因为节点经过平移,所以head(原始节点) -> head(平移节点) func (this *TreeChain) addSubtree(head int, value int) { // 原始点编号 -> 平移编号 head++ // 平移编号 -> dfs编号 dfn[head] this.seg.add(this.dfn[head], this.dfn[head]+this.siz[head]-1, value, 1, this.n, 1) } func (this *TreeChain) querySubtree(head int) int { head++ return this.seg.query(this.dfn[head], this.dfn[head]+this.siz[head]-1, 1, this.n, 1) } func (this *TreeChain) addChain(a int, b int, v int) { a++ b++ for this.top[a] != this.top[b] { if this.dep[this.top[a]] > this.dep[this.top[b]] { this.seg.add(this.dfn[this.top[a]], this.dfn[a], v, 1, this.n, 1) a = this.fa[this.top[a]] } else { this.seg.add(this.dfn[this.top[b]], this.dfn[b], v, 1, this.n, 1) b = this.fa[this.top[b]] } } if this.dep[a] > this.dep[b] { this.seg.add(this.dfn[b], this.dfn[a], v, 1, this.n, 1) } else { this.seg.add(this.dfn[a], this.dfn[b], v, 1, this.n, 1) } } func (this *TreeChain) queryChain(a int, b int) int { a++ b++ ans := 0 for this.top[a] != this.top[b] { if this.dep[this.top[a]] > this.dep[this.top[b]] { ans += this.seg.query(this.dfn[this.top[a]], this.dfn[a], 1, this.n, 1) a = this.fa[this.top[a]] } else { ans += this.seg.query(this.dfn[this.top[b]], this.dfn[b], 1, this.n, 1) b = this.fa[this.top[b]] } } if this.dep[a] > this.dep[b] { ans += this.seg.query(this.dfn[b], this.dfn[a], 1, this.n, 1) } else { ans += this.seg.query(this.dfn[a], this.dfn[b], 1, this.n, 1) } return ans } type SegmentTree struct { MAXN int arr []int sum []int lazy []int } func NewSegmentTree(origin []int) *SegmentTree { ret := &SegmentTree{} ret.MAXN = len(origin) ret.arr = origin ret.sum = make([]int, ret.MAXN<<2) ret.lazy = make([]int, ret.MAXN<<2) return ret } func (this *SegmentTree) pushUp(rt int) { this.sum[rt] = this.sum[rt<<1] + this.sum[rt<<1|1] } func (this *SegmentTree) pushDown(rt int, ln int, rn int) { if this.lazy[rt] != 0 { this.lazy[rt<<1] += this.lazy[rt] this.sum[rt<<1] += this.lazy[rt] * ln this.lazy[rt<<1|1] += this.lazy[rt] this.sum[rt<<1|1] += this.lazy[rt] * rn this.lazy[rt] = 0 } } func (this *SegmentTree) build(l int, r int, rt int) { if l == r { this.sum[rt] = this.arr[l] return } mid := (l + r) >> 1 this.build(l, mid, rt<<1) this.build(mid+1, r, rt<<1|1) this.pushUp(rt) } func (this *SegmentTree) add(L int, R int, C int, l int, r int, rt int) { if L <= l && r <= R { this.sum[rt] += C * (r - l + 1) this.lazy[rt] += C return } mid := (l + r) >> 1 this.pushDown(rt, mid-l+1, r-mid) if L <= mid { this.add(L, R, C, l, mid, rt<<1) } if R > mid { this.add(L, R, C, mid+1, r, rt<<1|1) } this.pushUp(rt) } func (this *SegmentTree) query(L int, R int, l int, r int, rt int) int { if L <= l && r <= R { return this.sum[rt] } mid := (l + r) >> 1 this.pushDown(rt, mid-l+1, r-mid) ans := 0 if L <= mid { ans += this.query(L, R, l, mid, rt<<1) } if R > mid { ans += this.query(L, R, mid+1, r, rt<<1|1) } return ans } // 为了测试,这个结构是暴力但正确的方法 type Right struct { n int tree [][]int fa []int val []int path map[int]int } func NewRight(father []int, value []int) *Right { ret := &Right{} ret.n = len(father) ret.tree = make([][]int, ret.n) ret.fa = make([]int, ret.n) ret.val = make([]int, ret.n) for i := 0; i < ret.n; i++ { ret.fa[i] = father[i] ret.val[i] = value[i] } help := make([]int, ret.n) h := 0 for i := 0; i < ret.n; i++ { if father[i] == i { h = i } else { help[father[i]]++ } } for i := 0; i < ret.n; i++ { ret.tree[i] = make([]int, help[i]) } for i := 0; i < ret.n; i++ { if i != h { help[father[i]]-- ret.tree[father[i]][help[father[i]]] = i } } ret.path = make(map[int]int) return ret } func (this *Right) addSubtree(head int, value int) { this.val[head] += value for _, next := range this.tree[head] { this.addSubtree(next, value) } } func (this *Right) querySubtree(head int) int { ans := this.val[head] for _, next := range this.tree[head] { ans += this.querySubtree(next) } return ans } func (this *Right) addChain(a int, b int, v int) { this.path = make(map[int]int) this.path[a] = math.MinInt64 for a != this.fa[a] { this.path[this.fa[a]] = a a = this.fa[a] } _, ok := this.path[b] for !ok { this.val[b] += v b = this.fa[b] } this.val[b] += v for this.path[b] != math.MaxInt64 { b = this.path[b] this.val[b] += v } } func (this *Right) queryChain(a int, b int) int { this.path = make(map[int]int) this.path[a] = math.MinInt64 for a != this.fa[a] { this.path[this.fa[a]] = a a = this.fa[a] } ans := 0 _, ok := this.path[b] for !ok { ans += this.val[b] b = this.fa[b] } ans += this.val[b] for this.path[b] != math.MinInt64 { b = this.path[b] ans += this.val[b] } return ans }