使用树状数组解决数组单点更新后快速查询区间和的问题

作者：Grey

原文地址：

博客园：使用树状数组解决数组单点更新后快速查询区间和的问题

CSDN：使用树状数组解决数组单点更新后快速查询区间和的问题

要解决的问题

在数组arr中，如何快速求数组区间[i...j]的累加和？

如果数组元素值不变，前缀和数组可以用来加速生成区间[i...j]的累加和信息, 方法如下：

假设前缀和数组为preSum，preSum[i]表示区间[0...i]的累加和，

那么区间[i...j]的累加和sum[i...j]可以通过如下公式获得：

sum[i...j] = preSum[j] - preSum[i-1]

前缀和数组preSum可以通过一次预处理即可获得，时间复杂度O(N),后续求任意区间的累加和只需要O(1)的时间复杂度。

前缀和的预处理方式如下：

preSum[0] = arr[0]
for (int i = 1 ; i < arr.length; i++) {
    preSum[i] = arr[i] + preSum[i-1];
}

示例：LeetCode 303. Range Sum Query - Immutable

但是如果数组会进行单点修改，例如：LeetCode 307. Range Sum Query - Mutable，则使用前缀和的方式来处理复杂度就比较高了，树状数组和线段树都可以用来解决这个问题，时间复杂度可以达到O(logN)。本文介绍树状数组的解法。

注：本文所有涉及到的数组操作均从下标1开始计算，下标0位置弃而不用

树状数组提供的方法

树状数组仅需要实现如下两个方法：

第一个方法：int sum(index)

表示区间[1...index]的累加和。

第二个方法：void add(index, d)

表示数组中index位置的值加d。

通过如上两个方法，

如果我们要求任意区间[left, right]的累加和。

只需要通过sum(right) - sum(left - 1)即可得到。

如果要将index位置的值更新为x。

只需要通过add(index, x - arr[index])即可得到。

初始代码如下：

// 树状数组
class NumArray {
    ....
    
    private void add(int index, int d) {
        ......
    }

    private int sum(int index) {
        ......
    }
}

预处理数组

为了方便处理，我们下标从1开始处理，所以假设原始数组为nums，数组长度为N，我们可以申请一个长度为N+1的辅助数组arr来存原始数组的信息，申请一个长度N+1的辅助数组tree来存累加和信息，tree数组一开始数据全部为0。然后开始遍历nums，填充arr和并根据一定的规则填充tree数组，当nums遍历完毕，arr数组和tree数组对应就生成好了。

代码如下：

// 树状数组
class NumArray {
    private int[] tree;
    // 原始数组长度
    private int N;
    private int[] arr;

    public NumArray(int[] nums) {
        N = nums.length;
        arr = new int[N + 1];
        tree = new int[N + 1];
        // 从1开始保存原始数组的信息
        System.arraycopy(nums, 0, arr, 1, N);
        for (int i = 1; i < tree.length; i++) {
            // i位置增加一个arr[i]的值
            add(i, arr[i]);
        }
    }
    
    private void add(int index, int d) {
        ......
    }

    private int sum(int index) {
        ......
    }
}

add方法

add方法表示：在index位置上的值增加一个d，此时需要考虑辅助数组tree哪些位置受到了牵连。规则如下：

从index位置开始，每次加上index最右侧的1，一直到数组结尾位置，都是受牵连的位置，这些位置都要执行加d操作。

class NumArray {
    ......
    private void add(int index, int d) {
        while (index <= N) {
            tree[index] += d;
            index += (index & (-index));
        }
    }
    ......
}

sum方法

按如上流程成tree数组后，如果要计算1..index位置的累加和，则有如下规则：

第一步，提取出index最右侧的1，假设为x，将help[index] + help[index - x]，得出a1
第二步，继续提取index-x最右侧的1，假设为y，将a1 + help[index - x - y]，得出a2
...
直到index提取完所有最右侧的1，求累加，得到的结果即为1...index上的累加和。

代码如下：

  
class NumArray {
    ......
    private int sum(int index) {
        int ret = 0;
        while (index > 0) {
            ret += tree[index];
            index -= (index & (-index));
        }
        return ret;
    }
    ......
}

树状数组完整代码

class NumArray {
    private int[] tree;
    private int N;
    private int[] arr;

    public NumArray(int[] nums) {
        arr = new int[nums.length + 1];
        System.arraycopy(nums, 0, arr, 1, nums.length);
        N = nums.length;
        tree = new int[N + 1];
        for (int i = 1; i < tree.length; i++) {
            add(i, arr[i]);
        }
    }
 
    public void add(int index, int d) {
        while (index <= N) {
            tree[index] += d;
            index += (index & (-index));
        }
    }
    
    public int sum(int index) {
        int ret = 0;
        while (index > 0) {
            ret += tree[index];
            index -= (index & (-index));
        }
        return ret;
    }
}

线段树 VS 树状数组

线段树是树状数组的升级版，树状数组只能做到单点更新后，维持累加和信息的快速更新，线段树可以支持范围更新，但是树状数组可以很方便改成二维或者三维的，对于线段树来说，改成二维的太复杂。

二维树状数组

二维树状数组主要解决：在单点变化时候，快速更新从左上角位置（1,1）累加到(i,j)位置的累加和信息这个问题。

熟悉一维数组后，二维的树状数组比较简单，原先一维数组只需要考虑1...i位置累加和，现在二维除了考虑1...i位置累加和，还要考虑1...j位置的累加和

二维树状数组的完整源码如下：

public class Code_0069_IndexTree2D {
    private int[][] tree;
    private int[][] nums;
    private int N;
    private int M;

    public Code_0069_IndexTree2D(int[][] matrix) {
        if (matrix.length == 0 || matrix[0].length == 0) {
            return;
        }
        N = matrix.length;
        M = matrix[0].length;
        tree = new int[N + 1][M + 1];
        nums = new int[N][M];
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < M; j++) {
                update(i, j, matrix[i][j]);
            }
        }
    }

    private int sum(int row, int col) {
        int sum = 0;
        for (int i = row + 1; i > 0; i -= i & (-i)) {
            for (int j = col + 1; j > 0; j -= j & (-j)) {
                sum += tree[i][j];
            }
        }
        return sum;
    }

    public void update(int row, int col, int val) {
        if (N == 0 || M == 0) {
            return;
        }
        int add = val - nums[row][col];
        nums[row][col] = val;
        for (int i = row + 1; i <= N; i += i & (-i)) {
            for (int j = col + 1; j <= M; j += j & (-j)) {
                tree[i][j] += add;
            }
        }
    }

    public int sumRegion(int row1, int col1, int row2, int col2) {
        if (N == 0 || M == 0) {
            return 0;
        }
        return sum(row2, col2) + sum(row1 - 1, col1 - 1) - sum(row1 - 1, col2) - sum(row2, col1 - 1);
    }
}

即在一维的条件下，增加了一个循环。现在有了二维树状数组，如果要求整个二维平面中，任意[row1,col1]位置到[row2,col2]位置组成的矩形累加和信息，则可以很方便通过二维数状数组计算出来，即代码中的sumRegion方法。