树状数组是一种动态维护前缀和的高效数据结构。它能在 O(log⁡n)的时间内完成单点更新前缀查询,空间复杂度 O(n),代码非常简洁。

普通数组求前缀和是 O(n)的,一旦需要频繁修改就太慢。树状数组把数组下标按照二进制分组,让每个节点管理一段区间和,使得查询和修改都只需跳少数几个节点。

1、单点修改

模板题:P3374 【模板】树状数组 1 - 洛谷

单点修改模板:

template<class T>
class FenwickTree
{
private:
	vector<T> m_tree;
    //最低位 1 对应的值
	T lowbit(T x);

public:
	FenwickTree(int n) :m_tree(n + 1, 0) {};
    //更新函数
	void Update(int idx, T val);
    //查询函数
	T Query(int idx);
	T Query(int l, int r);
};

template<class T>
inline T FenwickTree<T>::lowbit(T x) {
	return (-x) & x;
}

//以O(log⁡n)的时间复杂度去插入数,组当插入值改变时它的父节点的值也要改变
template<class T>
void FenwickTree<T>::Update(int idx, T val)
{
	int n = (int)m_tree.size();
	while (idx < n) {
		m_tree[idx] += val;
		idx += lowbit(idx);
	}
}

//以O(log⁡n)的时间复杂度获得前缀和
template<class T>
inline T FenwickTree<T>::Query(int idx) {
	T sun = 0;
	while (idx > 0) {
		sun += m_tree[idx];
		idx -= lowbit(idx);
	}
	return sun;
}

//查询单点时的值
template<class T>
inline T FenwickTree<T>::Query(int l, int r) {
	return Query(r) - Query(l - 1);
}

2、区间修改

模板题:P3368 【模板】树状数组 2 - 洛谷

区间修改模板:区间修改的代码就是将前缀和改为了差分,让我们在求前缀和的时候直接求得对应下标的值。

template<class T>
class FenwickTree
{
private:
	vector<T> m_tree;
	T lowbit(T x);
	void Update(int idx, T val);
	T Query(int idx);

public:
    //注意这里的vector容器的上限变为了n+2
	FenwickTree(int n) :m_tree(n + 2, 0) {};
    //区间修改
	void UpdateInterval(int l, int r, T idx);
    //直接查询下标值
	T QueryIndex(int idx);
};

template<class T>
inline T FenwickTree<T>::lowbit(T x) {
	return (-x) & x;
}

template<class T>
void FenwickTree<T>::Update(int idx, T val)
{
	int n = (int)m_tree.size();
	while (idx < n) {
		m_tree[idx] += val;
		idx += lowbit(idx);
	}
}

template<class T>
inline T FenwickTree<T>::Query(int idx) {
	T sun = 0;
	while (idx > 0) {
		sun += m_tree[idx];
		idx -= lowbit(idx);
	}
	return sun;
}

//区间修改
template<class T>
void FenwickTree<T>::UpdateInterval(int l, int r, T val)
{
	Update(l, val);
    //因为这里的r+1有可能会大于n所以容器大小变为n+2
	Update(r + 1, -val);
}

//直接查询下标值
template<class T>
inline T FenwickTree<T>::QueryIndex(int idx) {
	return Query(idx);
}