每日算法 — 使用java实现树状数组:单点修改与区间查询的高效数据结构

引言

树状数组(Binary Indexed Tree,简称BIT,又称Fenwick Tree)是一种极其优雅的数据结构,它能在 O(log n) 时间内完成单点修改和前缀和查询,同时只需要 O(n) 的额外空间。相比线段树,树状数组实现更简洁、常数更小,是解决区间查询与单点更新问题的首选工具。本文将用Java从零实现树状数组,深入讲解其数学原理 lowbit 运算,并展示其在逆序对计数、区间求和等经典问题中的应用。

核心思想

树状数组的核心是维护一个数组 tree[],其中每个元素 tree[i] 存储的是原数组中某一段区间的和。关键在于:

  • tree[i] 负责管理的区间长度为 lowbit(i),即 i 的二进制表示中最低位的1所代表的数值
  • tree[i] 存储的是区间 [i - lowbit(i) + 1, i] 的和

lowbit 运算:lowbit(x) = x & (-x),利用补码性质提取最低位的1。

例如:
– lowbit(6) = lowbit(110₂) = 2,所以 tree[6] 管理区间 [5, 6]
– lowbit(8) = lowbit(1000₂) = 8,所以 tree[8] 管理区间 [1, 8]

基础实现

/**
 * 树状数组(Binary Indexed Tree / Fenwick Tree)实现
 * 支持:单点修改、前缀和查询、区间和查询
 */
public class FenwickTree {

    private long[] tree;  // 树状数组(1-based)
    private int n;        // 数组大小

    /**
     * 构造函数:创建大小为n的树状数组
     */
    public FenwickTree(int n) {
        this.n = n;
        this.tree = new long[n + 1]; // 1-based索引
    }

    /**
     * 构造函数:从已有数组构建树状数组
     */
    public FenwickTree(long[] arr) {
        this.n = arr.length;
        this.tree = new long[n + 1];
        for (int i = 0; i < n; i++) {
            add(i + 1, arr[i]); // 逐个插入,O(n log n)
        }
    }

    /**
     * lowbit运算:提取最低位的1
     * 原理:x & (-x)
     * 例如:x = 12 = 1100₂, -x = ...11110100₂ (补码)
     *       x & (-x) = 0100₂ = 4
     */
    private int lowbit(int x) {
        return x & (-x);
    }

    /**
     * 单点修改:在位置index增加delta
     * 时间复杂度: O(log n)
     *
     * 修改路径:index -> index + lowbit(index) -> ...
     * 这些节点都包含了index位置的信息,因此都需要更新
     */
    public void add(int index, long delta) {
        // 注意:外部传入的index是1-based
        while (index <= n) {
            tree[index] += delta;
            index += lowbit(index); // 跳到父节点
        }
    }

    /**
     * 前缀和查询:查询[1, index]的和
     * 时间复杂度: O(log n)
     *
     * 查询路径:index -> index - lowbit(index) -> ... -> 0
     * 每次减去最低位的1,将区间分解为若干段不相交的区间之和
     */
    public long prefixSum(int index) {
        long sum = 0;
        while (index > 0) {
            sum += tree[index];
            index -= lowbit(index); // 跳到前一个区间
        }
        return sum;
    }

    /**
     * 区间和查询:查询[left, right]的和
     * 利用前缀和:sum(left, right) = prefix(right) - prefix(left-1)
     */
    public long rangeSum(int left, int right) {
        if (left > right) return 0;
        return prefixSum(right) - prefixSum(left - 1);
    }

    /**
     * 单点查询:查询位置index的值
     * 等价于 rangeSum(index, index)
     */
    public long pointQuery(int index) {
        return rangeSum(index, index);
    }

    /**
     * 单点修改(设置值):将位置index设置为value
     */
    public void set(int index, long value) {
        long current = pointQuery(index);
        add(index, value - current); // 增加差值
    }

    /**
     * 打印树状数组结构(调试用)
     */
    public void printStructure() {
        System.out.println("=== 树状数组结构 ===");
        for (int i = 1; i <= n; i++) {
            int lb = lowbit(i);
            System.out.println("tree[" + i + "] = " + tree[i] +
                               " (lowbit=" + lb + ", 管理区间 [" + (i - lb + 1) + ", " + i + "])");
        }
    }

数学原理可视化

    /**
     * 展示树状数组的区间分解原理
     */
    public static void demonstrateDecomposition() {
        System.out.println("=== 树状数组区间分解原理 ===");
        System.out.println("前缀和 query(13) 的分解:");
        int x = 13; // 1101₂
        System.out.println("13 = " + Integer.toBinaryString(x) + "₂");

        // 分解过程
        int step = 1;
        while (x > 0) {
            int lb = x & (-x);
            System.out.println("  步骤" + step + ": x=" + x + ", lowbit=" + lb +
                               ", 取 tree[" + x + "](管理区间 [" + (x - lb + 1) + ", " + x + "])");
            x -= lb;
            step++;
        }
        System.out.println("  最终:prefix(13) = tree[13] + tree[12] + tree[8]");
        System.out.println("  对应区间:[13,13] + [9,12] + [1,8] = [1,13]");
        System.out.println();

        System.out.println("修改 update(5, +3) 的影响路径:");
        x = 5; // 0101₂
        System.out.println("5 = " + Integer.toBinaryString(x) + "₂");
        while (x <= 16) {
            int lb = x & (-x);
            System.out.println("  更新 tree[" + x + "](管理区间 [" + (x - lb + 1) + ", " + x + "])");
            x += lb;
        }
        System.out.println("  最终:tree[5], tree[6], tree[8], tree[16] 都被更新");
    }

经典应用一:逆序对计数

逆序对是数组中满足 i < j 且 a[i] > a[j] 的元素对。树状数组可在 O(n log n) 内高效计数。

    /**
     * 应用1:逆序对计数
     * 思路:从右向左遍历,查询已遍历元素中比当前元素小的个数
     *
     * @param arr 输入数组
     * @return 逆序对总数
     */
    public static long countInversions(int[] arr) {
        int n = arr.length;

        // 离散化:将数值映射到排名(1-based)
        // 这里简化处理,假设数值范围在 [1, MAX] 内
        // 实际应用需要排序去重后映射
        int maxVal = 0;
        for (int v : arr) maxVal = Math.max(maxVal, v);

        FenwickTree bit = new FenwickTree(maxVal);
        long inversions = 0;

        // 从右向左遍历
        for (int i = n - 1; i >= 0; i--) {
            // 查询已经遍历过的元素中,值小于 arr[i] 的个数
            // 这些元素在原数组中位于 arr[i] 右侧且值更小,构成逆序对
            inversions += bit.prefixSum(arr[i] - 1);
            // 将当前元素加入树状数组
            bit.add(arr[i], 1);
        }

        return inversions;
    }

    /**
     * 带离散化的逆序对计数(处理大范围数值)
     */
    public static long countInversionsWithDiscretization(int[] arr) {
        int n = arr.length;

        // 离散化:复制、排序、去重、映射
        int[] sorted = arr.clone();
        java.util.Arrays.sort(sorted);

        // 去重后的唯一值列表(二分查找用)
        int m = 0;
        for (int i = 0; i < n; i++) {
            if (i == 0 || sorted[i] != sorted[i - 1]) {
                sorted[m++] = sorted[i];
            }
        }

        FenwickTree bit = new FenwickTree(m);
        long inversions = 0;

        for (int i = n - 1; i >= 0; i--) {
            // 二分查找排名
            int rank = java.util.Arrays.binarySearch(sorted, 0, m, arr[i]) + 1;
            inversions += bit.prefixSum(rank - 1);
            bit.add(rank, 1);
        }

        return inversions;
    }

经典应用二:区间最值与二维树状数组

    /**
     * 应用2:二维树状数组
     * 支持矩阵的单点修改和子矩阵求和
     */
    static class FenwickTree2D {
        private long[][] tree;
        private int n, m;

        FenwickTree2D(int n, int m) {
            this.n = n;
            this.m = m;
            this.tree = new long[n + 1][m + 1];
        }

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

        /**
         * 单点修改:在(x, y)位置增加delta
         */
        public void add(int x, int y, long delta) {
            for (int i = x; i <= n; i += lowbit(i)) {
                for (int j = y; j <= m; j += lowbit(j)) {
                    tree[i][j] += delta;
                }
            }
        }

        /**
         * 前缀和查询:[1,1] 到 [x,y] 的矩形区域和
         */
        public long prefixSum(int x, int y) {
            long sum = 0;
            for (int i = x; i > 0; i -= lowbit(i)) {
                for (int j = y; j > 0; j -= lowbit(j)) {
                    sum += tree[i][j];
                }
            }
            return sum;
        }

        /**
         * 子矩阵和查询:[x1,y1] 到 [x2,y2]
         */
        public long rangeSum(int x1, int y1, int x2, int y2) {
            return prefixSum(x2, y2) - prefixSum(x1 - 1, y2)
                 - prefixSum(x2, y1 - 1) + prefixSum(x1 - 1, y1 - 1);
        }
    }

完整测试与验证

    /**
     * 主程序:测试树状数组的所有功能
     */
    public static void main(String[] args) {
        // 演示lowbit原理
        demonstrateDecomposition();
        System.out.println();

        // 基础功能测试
        System.out.println("=== 基础功能测试 ===");
        long[] arr = {0, 1, 7, 3, 0, 5, 8, 3, 2, 6}; // 0号位置不用,1-based
        FenwickTree bit = new FenwickTree(9);

        // 初始化
        for (int i = 1; i <= 9; i++) {
            bit.add(i, arr[i]);
        }

        bit.printStructure();
        System.out.println();

        // 前缀和查询
        System.out.println("前缀和测试:");
        for (int i = 1; i <= 9; i++) {
            System.out.println("prefixSum(" + i + ") = " + bit.prefixSum(i));
        }
        System.out.println();

        // 区间和查询
        System.out.println("区间和测试:");
        System.out.println("rangeSum(3, 7) = " + bit.rangeSum(3, 7)); // 3+0+5+8+3 = 19
        System.out.println("rangeSum(1, 9) = " + bit.rangeSum(1, 9)); // 总和 = 35
        System.out.println();

        // 单点修改测试
        System.out.println("单点修改测试: add(4, 5)");
        bit.add(4, 5); // 位置4增加5
        System.out.println("修改后 prefixSum(4) = " + bit.prefixSum(4));
        System.out.println("修改后 rangeSum(3, 5) = " + bit.rangeSum(3, 5));
        System.out.println();

        // 逆序对计数测试
        System.out.println("=== 逆序对计数测试 ===");
        int[] testArr1 = {5, 2, 6, 1, 3};
        System.out.println("数组: " + java.util.Arrays.toString(testArr1));
        System.out.println("逆序对数: " + countInversions(testArr1));
        // 逆序对:(5,2), (5,1), (5,3), (2,1), (6,1), (6,3) = 6个

        int[] testArr2 = {1, 2, 3, 4, 5};
        System.out.println("数组: " + java.util.Arrays.toString(testArr2));
        System.out.println("逆序对数: " + countInversions(testArr2)); // 0个

        int[] testArr3 = {5, 4, 3, 2, 1};
        System.out.println("数组: " + java.util.Arrays.toString(testArr3));
        System.out.println("逆序对数: " + countInversions(testArr3)); // 10个
        System.out.println();

        // 二维树状数组测试
        System.out.println("=== 二维树状数组测试 ===");
        FenwickTree2D bit2d = new FenwickTree2D(4, 4);
        // 初始化矩阵
        int[][] matrix = {
            {0, 0, 0, 0},
            {0, 1, 2, 3},
            {0, 4, 5, 6},
            {0, 7, 8, 9}
        };
        for (int i = 1; i <= 3; i++) {
            for (int j = 1; j <= 3; j++) {
                bit2d.add(i, j, matrix[i][j]);
            }
        }
        System.out.println("子矩阵 [2,2] 到 [3,3] 的和 = " + bit2d.rangeSum(2, 2, 3, 3));
        // 5 + 6 + 8 + 9 = 28
        System.out.println("子矩阵 [1,1] 到 [3,3] 的和 = " + bit2d.rangeSum(1, 1, 3, 3));
        // 1+2+3+4+5+6+7+8+9 = 45
        System.out.println();

        // 性能测试
        System.out.println("=== 性能测试 ===");
        int size = 1000000;
        FenwickTree bigBit = new FenwickTree(size);

        long start = System.currentTimeMillis();
        for (int i = 1; i <= size; i++) {
            bigBit.add(i, i);
        }
        long buildTime = System.currentTimeMillis() - start;

        start = System.currentTimeMillis();
        long sum = bigBit.prefixSum(size);
        long queryTime = System.currentTimeMillis() - start;

        System.out.println("规模: " + size);
        System.out.println("构建时间: " + buildTime + "ms");
        System.out.println("前缀和查询时间: " + queryTime + "ms");
        System.out.println("前缀和结果: " + sum);
        System.out.println("期望结果: " + ((long) size * (size + 1) / 2));
    }
}

复杂度分析

操作 时间复杂度 空间复杂度 说明
单点修改 add O(log n) O(1) 沿父节点路径更新
前缀和查询 O(log n) O(1) 区间分解累加
区间和查询 O(log n) O(1) 两次前缀和差分
单点赋值 set O(log n) O(1) 先查询再更新差值
初始化构建 O(n log n) O(n) 逐个插入;可优化到O(n)
逆序对计数 O(n log n) O(n) 需要离散化

构建优化:通过直接计算 tree[i] = prefixSum(i) - prefixSum(i - lowbit(i)) 可在 O(n) 内完成构建。

树状数组 vs 线段树

特性 树状数组 线段树
代码量 短(约30行) 较长(约100+行)
常数因子 小 较大
单点修改+区间求和 ✓ ✓
区间修改+单点查询 ✓(差分思想) ✓(懒标记)
区间修改+区间查询 需配合差分,较复杂 ✓(懒标记)
区间最值 ✗ 不支持 ✓
可扩展性 有限 强

扩展应用

  1. 区间修改 + 单点查询:维护差分数组,用树状数组支持区间加、单点查
  2. 区间修改 + 区间查询:维护两个树状数组,分别记录差分和差分前缀和
  3. 动态排名查询:结合二分查找,在 O(log² n) 内找到第 k 大元素
  4. 多维树状数组:扩展到三维乃至更高维,处理空间立方体区域求和

总结

树状数组以其极简的实现和优秀的性能成为算法竞赛和工程实践中的常客。其核心 lowbit 运算巧妙地利用了二进制性质,将数组组织成一棵隐式的树结构。无论是逆序对计数、区间求和还是动态排名问题,树状数组都能以 O(log n) 的时间高效解决。对于只需要单点修改和区间求和的场景,树状数组是线段树的轻量级完美替代。