树状数组理论与实践

1 理论

1.1 介绍

树状数组(Fenwick Tree)，也称为二叉索引树（Binary Indexed Tree, BIT），是一种用于高效计算数组前缀和的数据结构。它可以在$O(logN)$的时间复杂度下完成单点更新、前缀和查询等操作。

树状数组的主要思想是利用二进制的特性来进行有效的更新和查询。它的核心思想是利用数组的索引来表示树状结构，从而实现高效的查询和更新操作。

具体来说，树状数组的结构如下：

1. 数组C：这是一个额外的数组，用于存储原始数组的前缀和或其它需要维护的信息。
2. 每个元素的索引i代表了原始数组中的某个位置。
3. 每个元素的值C[i]存储了原始数组中某个位置到其父节点（通过某种规则计算）的信息。通常，C[i]的计算规则是将i的二进制表示中最低位的1所代表的区间累加到i上。

1.2 lowbit操作

lowbit(x)表示数x二进制表示中最后一位1的大小，用十进制表示。

int lowbit(int x) {
	return x & -x;
}

例如输入6，返回2；输入7返回1；输入8返回8；输入9返回1。

对于数x，加上lowbit(x)的写法，

//写法1
x += lowbit(x);
//写法2
x += x & -x;

对于数x，减去lowbit(x)的写法，

//写法1
x -= lowbit(x);
//写法2
x -= x & -x;
//写法3
x &= x - 1;

1.3 C++模板

树状数组的C++模板如下，

//给定数组nums: 下标从0开始
//(1)构建tree数组
int n = nums.size();
vector<int> tree(n + 1, 0); //下标从1开始，故是n + 1。nums[i] <==> tree[i + 1]。
//tree[k]表示以k为区间右端点的那个区间下所有元素的和，也即[k - lowbit(k) + 1, k]，该区间有lowbit(k)个元素。根据“tree数组从下标1开始，nums数组从下标0开始的设定”，转换成nums有：nums下标在[k - lowbit(k), k - 1]内的所有元素的和，该区间内有lowbit(k)个元素。

//(2.1)nums[i]增加操作
//nums[i] -> nums[i] + x
//nums[i]增加x，对相应的关键区间tree[k]的影响
void add(int i, int x) {
	nums[i] += x; //nums[i]增加x
	for (int k = i + 1; k < tree.size(); k += lowbit(k)) {
		tree[k] += x;
	}
	return;
}

//(2.2)nums[i]更新操作
//nums[i] -> val
//nums[i]更新为val，对相应的关键区间tree[k]的影响
void update(int i, int val) {
	int x = val - nums[i];
	add(i, x);
	return;
}

//(3)初始化
for (int i = 0; i < n; ++i) {
	add(i, nums[i]);
}

//(4.1)求nums数组中[0, r]区间内的元素和。即nums[0] + nums[1] + ... + nums[r]的和。
int get_s(int r) {
	int res = 0;
	for (int k = r + 1; k > 0; k -= lowbit(k)) {
		res + tree[k];
	}
	return res;
}

//(4.2)求nums数组中[l, r]区间内的元素和。即nums[l] + nums[l + 1] + ... + nums[r]的和。
int get_s(int l, int r) {
	return get_s(r) - get_s(l - 1);
}

重点：
tree[k]：以k为区间右端点的那个区间下所有元素的和，也即[k - lowbit(k) + 1, k]，该区间有lowbit(k)个元素。根据“tree数组从下标1开始，nums数组从下标0开始的设定”，转换成nums有：nums下标在[k - lowbit(k), k - 1]内的所有元素的和，该区间内有lowbit(k)个元素。

2 实践

模板题：307区域和检索-数组可修改

C++代码如下，

class NumArray {
private:
    vector<int> nums;
    vector<int> tree;

public:
    inline int lowbit(int x) {
        return x & -x;
    }

    NumArray(vector<int>& nums) {
        int n = nums.size();
        this->nums.resize(n, 0);
        tree.resize(n + 1, 0);

        for (int i = 0; i < n; ++i) {
            update(i, nums[i]);
        }

    }
    
    void update(int index, int val) {
        int x = val - nums[index];

        for (int k = index + 1; k < tree.size(); k += lowbit(k)) {
            tree[k] += x;
        }

        nums[index] = val;//更新nums[index]

        return;
    }
    
    int sumRange(int right) {
        int res = 0;
        for (int k = right + 1; k > 0; k -= lowbit(k)) {
            res += tree[k];
        }
        return res;
    }

    int sumRange(int left, int right) {
        return sumRange(right) - sumRange(left - 1);
    }
};

题目1：3072. 将元素分配到两个数组中 II

C++代码如下，

class Fenwick {
    vector<int> tree;
public:
    Fenwick(int n) : tree(n) {} //下标从1开始

    //把下标为i的元素增加1
    void add(int i) {
        while (i < tree.size()) {
            tree[i]++;
            i += i & -i;
        }
    }

    //返回下标在[1,i]的元素之和,注意是闭区间
    int pre(int i) {
        int res = 0;
        while (i > 0) {
            res += tree[i];
            i &= i - 1;
        }
        return res;
    }
};

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        auto sorted = nums;
        ranges::sort(sorted);
        sorted.erase(unique(sorted.begin(), sorted.end()), sorted.end());
        int m = sorted.size();

        vector<int> a{nums[0]}, b{nums[1]};
        Fenwick t1(m + 1), t2(m + 1); //t1,t2是树状数组
        t1.add(ranges::lower_bound(sorted, nums[0]) - sorted.begin() + 1); //+1是因为下标从1开始
        t2.add(ranges::lower_bound(sorted, nums[1]) - sorted.begin() + 1); //+1是因为下标从1开始
        for (int i = 2; i < nums.size(); ++i) {
            int x = nums[i];
            int v = ranges::lower_bound(sorted, x) - sorted.begin() + 1;
            int gc1 = a.size() - t1.pre(v);
            int gc2 = b.size() - t2.pre(v);
            if (gc1 > gc2 || gc1 == gc2 && a.size() <= b.size()) {
                a.push_back(x);
                t1.add(v);
            } else {
                b.push_back(x);
                t2.add(v);
            }
        }
        a.insert(a.end(), b.begin(), b.end());
        return a;
    }
};

题目2：

3 参考

1. 带你发明树状数组！附数学证明（Python/Java/C++/Go/JS/Rust）