树状数组（速食版）

2024-06-06

字数统计: 1.4k字 | 阅读时长: 6分

前言

快速计算任意连续子数组元素和的数据结构。

树状数组或二元索引树（英语：Binary Indexed Tree），又以其发明者命名为
Fenwick树。最早由 PeterM.Fenwick于1994年以《A New Data Structure for Cumulative Frequency Tables》为题发表在《SOFTWARE PRACTICE AND EXPERIENCE》。

树状数组

给你一个数组，如何快速地计算任意一段连续子数组的元素和？

bit

起源

下标从 $left$ 到 $right$ 的子数组元素和，可以看成是下标从 1 到 $right$ 的子数组元素和，减去下标从 1 到 $left−1$ 的子数组元素和。例如数组 [3,1,4,1,5,9]，子数组 [4,1,5] 的元素和，等于 [3,1,4,1,5] 的元素和，减去 [3,1] 的元素和。

按照这个方法，算出每个前缀 [1,i]（表示下标从 1 到 i 的连续子数组）的元素和，就可以 $O(1)$ 地查询（计算任意连续子数组的元素和）了。

如果更新呢？

如果修改下标为1的元素，则所有前缀都需要更新（因为所有前缀包含该元素），意味着更新操作时间复杂度为$O(n)$。此时查询与更新的综合时间复杂度还是$O(n)$。

自然的，想到应该将前缀元素和继续进行拆分，令更新某元素时，仅影响部分“子前缀”，此时只需要更新影响到的部分即可。

因此，需要找到一种合适的拆分方法，能够把任意前缀拆分成若干子前缀，使更新操作的执行范围缩小到部分子前缀中。

对于想继续探索的同学

推荐文献[3]，较为全面的阐述了树状数组的来世今生，本文以让读者尽快掌握用法为主，不过多展开深入。

lowbit

BIT的重要概念，抛开相关性质不谈，先说计算，lowbit(x)表示取数x的最低位1，常用lowbit(x) = x & -x计算。

举例：

x = 1010, lowbit(x) = x & -x = 10

计算过程：
   x = 1010
  -x = 0110
x&-x = 0010

树状数组 `c[n]` 的查询与更新

对于原数组a[n]构建对应的树状数组c[n]，可推导得出：c[x] = a[x-lowbit(x)+1] + ... + a[x]，即区间[x-lowbit(x)+1, x]的元素和。

定义query(x)为查询原数组区间[1, x]上的元素和，则可写出利用c[n]实现$O(logn)$的查询代码：

int query(int x)
{
    int ans = 0;
    // 减去所有的lowbit后i为0
    for(int i = x; i ; i -= lowbit(i))
            ans += c[i];
    return ans;
}

或者是：

int query(int x)
{
    int ans = 0;
    while (x > 0){
        ans += c[x];
        x -= lowbit(x);
    }
    return ans;
}

定义add(x, v)为修改位置为在索引x的元素，加上v。则可写出利用c[n]实现$O(logn)$的更新代码：

// n = 树状数组长度
void add(int x, int v)
{
    while (x <= n){
        c[x] += v;
        x += lowbit(x);
    }
}

以上内容很好理解，抓住c[x]的定义就行了。

树状数组 `c[n]` 的构造

最简单的，遍历原数组，调用n次add方法进行构造，时间复杂度$O(nlogn)$：

void init()
{
    for(int i = 1; i <= n; i++){
        add(i, a[i]);
    }
}

更快速的，考虑c[x]表示原数组区间[x-lowbit(x)+1, x]的元素和，可以先对原数组求一个前缀和数组s[n]，利用前缀和来更新树状数组，即c[x] = s[x] - s[x-lowbit(x)]，此时构造的时间复杂度为$O(n)$：

void init()
{
    for(int i = 1; i <= n; i ++) // 求a的前缀和
        s[i] = s[i - 1] + a[i];
    for(int i = 1; i <= n; i ++) // 用前缀和求出c
		c[i] = s[i] - s[i - lowbit(i)];
}

拆分规则—逆序

按照逆序处理，每次处理的bit都是当前编号的最后的为1位。将每次处理的bit定义为lowbit。

对于前缀 $[1,i]$：

如果 $i$ 是 2 的幂，那么 $[1,i]$ 无需拆分。
如果 $i$ 不是 2 的幂，那么先拆分出一个长为 $lowbit(i)$ 的关键区间 $[i−lowbit(i)+1,i]$，问题转换成剩下的 $[1,i−lowbit(i)]$ 如何拆分，这是一个规模更小的子问题。

例题一：分数字到两个数组

class BinaryIndexedTree {
    private int[] tree;

    public int lowbit(int i){
        return i & -i;
    }

    public BinaryIndexedTree(int n) {
        tree = new int[n + 1];
    }

    public void add(int i) {
        while (i < tree.length) {
            tree[i]++;
            i += lowbit(i);
        }
    }

    public int get(int i) {
        int sum = 0;
        while (i > 0) {
            sum += tree[i];
            i -= lowbit(i);
        }
        return sum;
    }
}

class Solution {
    public int[] resultArray(int[] nums) {
        int n = nums.length;
        int[] sortedNums = Arrays.copyOf(nums, n);
        Arrays.sort(sortedNums);

        Map<Integer, Integer> index = new HashMap<>();
        for (int i = 0; i < n; i++) {
            index.put(sortedNums[i], i + 1);
        }

        List<Integer> arr1 = new ArrayList<>(List.of(nums[0]));
        List<Integer> arr2 = new ArrayList<>(List.of(nums[1]));
        BinaryIndexedTree tree1 = new BinaryIndexedTree(n);
        BinaryIndexedTree tree2 = new BinaryIndexedTree(n);
        tree1.add(index.get(nums[0]));
        tree2.add(index.get(nums[1]));

        for (int i = 2; i < n; i++) {
            int count1 = arr1.size() - tree1.get(index.get(nums[i]));
            int count2 = arr2.size() - tree2.get(index.get(nums[i]));
            if (count1 > count2 || (count1 == count2 && arr1.size() <= arr2.size())) {
                arr1.add(nums[i]);
                tree1.add(index.get(nums[i]));
            } else {
                arr2.add(nums[i]);
                tree2.add(index.get(nums[i]));
            }
        }

        int i = 0;
        for (int a: arr1) {
            nums[i++] = a;
        }
        for (int a: arr2) {
            nums[i++] = a;
        }
        return nums;
    }
}

参考文献

[1] https://leetcode.cn/problems/range-sum-query-mutable/solutions/2524481/dai-ni-fa-ming-shu-zhuang-shu-zu-fu-shu-lyfll/
[2] https://blog.csdn.net/qq_63786973/article/details/127416700
[3] https://www.cnblogs.com/Last--Whisper/p/13823614.html

前言