数据结构-树状数组

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

class TreeArray {
    private int[] tree; // sum(nums[i])
    private int n;

    public TreeArray(int[] nums) {
        n = nums.length + 1;
        tree = new int[n];
        for (int i = 0; i < n - 1; i++) {
            add(i, nums[i]);
        }
    }

    public void add(int index, int val) {
        // 下标从 1 开始
        index++;
        // 单点修改，增加数组 index 元素的值
        while (index < n) {
            tree[index] += val;
            // 更新父结点
            index += lowBit(index);
        }
    }

    public int preSum(int index) {
        // 查询前缀和
        int sum = 0;
        while (index > 0) {
            sum += tree[index];
            // 查询子结点
            index -= lowBit(index);
        }
        return sum;
    }

    private static int lowBit(int x) {
        // 返回 x 二进制最低位 1 的值
        // eg. 6(0b110) 返回 2(0b010)
        return x & (-x);
    }
}