[go]树状数组的实现

这个up主讲的非常细,我这里只是简单实现和使用

package treenums

type TreeNum struct {
	nums []int
}

// NewTreeNum 初始化一个存储n个元素的树状数组
func NewTreeNum(n int) *TreeNum {
	return &TreeNum{nums: make([]int, n+1)}
}

//计算x二进制的最低位1及之后的数。101000 -> 001000
func (t *TreeNum) lowbit(x int) int {
	return x & (-x)
}

// Add x点值+d;x = [0,n-1]
func (t *TreeNum) Add(x, d int) {
	x++
	for ; x < len(t.nums); x += t.lowbit(x) { //因为动态维护一个前缀和数组，所以依次都要+d
		t.nums[x] += d
	}
}

// Ask 查询x下标处的前缀和;x = [0,n-1]
func (t *TreeNum) Ask(x int) (ret int) {
	x++
	for ; x >= 1 && x < len(t.nums); x -= t.lowbit(x) {
		ret += t.nums[x]
	}
	return
}

/*
单点修改，区间查询
	[x]+d == add(x,d)
	[l,r] == ask(r)-ask(l-1)
区间修改，单点查询(使用树状数组维护一个查分数组b，记录查分数组的前缀和)
	[l,r]+d == add(l,d);add(r+1,-d) //[l]+d使得后面的都+d，[r+1]-d 保证区间外的值不受影响
	a[x] == ask(x)+a[x] // 查分数组记录了区间的值得变化情况
*/

树状数组用来维护前缀和非常好用

1. 307. 区域和检索 - 数组可修改

type NumArray struct {
	*TreeNum
	nums []int
}

func Constructor(nums []int) NumArray {
	treeNum := NewTreeNum(len(nums))
	for i := range nums {
		treeNum.add(i, nums[i])
	}
	return NumArray{
		TreeNum: treeNum,
		nums:    nums,
	}
}
//更新某一点的值
func (this *NumArray) Update(index int, val int) {
	a := this.nums[index]
	this.add(index, val-a)
	this.nums[index] = val
}
//计算区间和
func (this *NumArray) SumRange(left int, right int) int {
	return this.ask(right) - this.ask(left-1)
}

type TreeNum struct {
	nums []int
}

func NewTreeNum(n int) *TreeNum {
	return &TreeNum{make([]int, n+1)} //下标从1开始
}

//计算x二进制的最低位1及之后的数。101000 -> 001000
func (t *TreeNum) lowbit(x int) int {
	return x & (-x)
}

//x点值+d
func (t *TreeNum) add(x, d int) {
	x++
	for ; x < len(t.nums); x += t.lowbit(x) { //因为动态维护一个前缀和数组，所以依次都要+d
		t.nums[x] += d
	}
}

//查询x下标处的前缀和
func (t *TreeNum) ask(x int) (ret int) {
	x++
	for ; x >= 1 && x < len(t.nums); x -= t.lowbit(x) {
		ret += t.nums[x]
	}
	return
}

2. 面试题 10.10. 数字流的秩

type StreamRank struct {
	*TreeNum
}

func Constructor() StreamRank {
	max := 50001
	return StreamRank{
		TreeNum: NewTreeNum(max),
	}
}
//修改值
func (this *StreamRank) Track(x int) {
	this.add(x, 1)
}
//获取点
func (this *StreamRank) GetRankOfNumber(x int) int {
	return this.ask(x)
}

type TreeNum struct {
	nums []int
}

func NewTreeNum(n int) *TreeNum {
	return &TreeNum{make([]int, n+1)} //下标从1开始
}

//计算x二进制的最低位1及之后的数。101000 -> 001000
func (t *TreeNum) lowbit(x int) int {
	return x & (-x)
}

//x点值+d
func (t *TreeNum) add(x, d int) {
	x++
	for ; x < len(t.nums); x += t.lowbit(x) { //因为动态维护一个前缀和数组，所以依次都要+d
		t.nums[x] += d
	}
}

//查询x下标处的前缀和
func (t *TreeNum) ask(x int) (ret int) {
	x++
	for ; x >= 1 && x < len(t.nums); x -= t.lowbit(x) {
		ret += t.nums[x]
	}
	return
}