[go刷题模板] 树状数组

一、 算法&数据结构

1. 描述

PURQ的树状数组可以用来单点更新区间查询

• 模板慢慢补充，更详细的可以看我的py版本。

2. 复杂度分析

1. 查询query, O(log₂n)
2. 更新update，O(log₂n)

3. 常见应用

4. 常用优化

二、 模板代码

1. 单点更新，区间询问求和(PURQ)

• cf652d D. Nested Segments
  

package main

import (
	"bufio"
	. "fmt"
	"io"
	"os"
	"sort"
)

// BIT 420 ms
func solve(in io.Reader, out io.Writer) {
	var n int
	Fscan(in, &n)
	type seg struct{ i, l, r int }
	var hs []int
	a := make([]seg, n)
	for i := 0; i < n; i++ {
		Fscan(in, &a[i].l, &a[i].r)
		hs = append(hs, []int{a[i].l, a[i].r}...)
		a[i].i = i
	}
	sort.Slice(a, func(i, j int) bool { return a[i].r < a[j].r })
	sort.Ints(hs)
	ans := make([]int, n)
	p := NewBIT(n * 2)
	for _, s := range a {
		l, i := s.l, s.i
		l = sort.Search(len(hs), func(i int) bool { return hs[i] >= l })
		ans[i] = p.SumInterval(l+1, n*2+1)
		p.AddPoint(l+1, 1)
	}

	for _, v := range ans {
		Fprintln(out, v)
	}
}
func run(_r io.Reader, _w io.Writer) {
	in := bufio.NewReader(_r)
	out := bufio.NewWriter(_w)
	defer out.Flush()
	solve(in, out)
}

func main() { run(os.Stdin, os.Stdout) }


type BIT struct {
	C []int
	N int
}

func NewBIT(n int) BIT {
	c := make([]int, n+2)
	return BIT{c, n}
}

// AddPoint
// 给i位置增加v，注意i是1-indexed
func (u *BIT) AddPoint(i, v int) {
	for i <= u.N {
		u.C[i] += v
		i += i & -i
	}
}

// SumPrefix
// 计算[..i]的前缀和，包含i，注意i是1-indexed
func (u BIT) SumPrefix(i int) int {
	s := 0
	for i >= 1 {
		s += u.C[i]
		i &= i - 1
	}
	return s
}

// SumInterval SumIn
// 计算[l..r]闭区间的和，注意l\r是1-indexed
func (u BIT) SumInterval(l, r int) int {
	return u.SumPrefix(r) - u.SumPrefix(l-1)
}

// Lowbit
// 计算lowbit，L
func (u BIT) Lowbit(x int) int {
	return x & -x
}

// MinRight 查找[i..]后第一个>=0的位置，包含i，注意i是1-indexed
func (u BIT) MinRight(i int) int {
	p := u.SumPrefix(i)
	if i == 1 {
		if p > 0 {
			return i
		}
	} else {
		if p > u.SumPrefix(i-1) {
			return i
		}
	}
	l, r := i, u.N+1
	for l+1 < r {
		mid := (l + r) >> 1
		if u.SumPrefix(mid) > p {
			r = mid
		} else {
			l = mid
		}
	}
	return r
}

三、其他

四、更多例题

五、参考链接