每日算法 — 使用java实现树状数组:二进制索引与区间查询的O(log n)解法

引言

在竞赛编程和工程实践中,频繁遇到这样的需求:维护一个数组,支持快速修改某个元素的值,以及快速查询某个区间的和。如果使用普通数组,修改是 O(1) 但区间查询是 O(n);如果使用前缀和数组,查询是 O(1) 但修改是 O(n)。树状数组(Binary Indexed Tree / Fenwick Tree)由 Peter Fenwick 于1994年提出,它通过巧妙的二进制索引设计,在 O(log n) 时间内同时支持单点修改和前缀和查询,空间复杂度仅为 O(n)。它是线段树的轻量级替代品,在只需维护区间和(或满足特定结合律的操作)的场景下,代码更短、常数更小。

本文将用 Java 从零实现树状数组,深入讲解其 lowbit 原理、区间查询推导、二维扩展以及逆序对统计等经典应用。

核心思想

前缀和的分解

树状数组的核心洞察在于:任何整数都可以分解为若干个 2 的幂次之和,对应地,任何前缀和也可以分解为若干段已知区间之和。

例如:
– 13 = 1101₂ = 8 + 4 + 1
– 前缀和 prefix[13] = tree[13] + tree[12] + tree[8]

其中 tree[i] 负责管理一段长度为 lowbit(i) 的区间。

lowbit 运算

lowbit(x) = x & (-x),返回 x 的二进制表示中最低位的 1 所对应的值。

x 二进制 -x (补码) x & (-x)
6 0110 1010 0010 = 2
8 1000 1000 1000 = 8
13 1101 0011 0001 = 1

树状数组的管辖区间

tree[i] 存储的是区间 [i – lowbit(i) + 1, i] 的元素之和,即以 i 结尾、长度为 lowbit(i) 的区间和。

Java 实现

基础一维树状数组

public class FenwickTree {
    private final int n;       // 数组长度
    private final long[] tree; // 树状数组内部存储(1-based)

    public FenwickTree(int n) {
        this.n = n;
        this.tree = new long[n + 1]; // 使用 1-based 索引
    }

    /**
     * 计算 lowbit:提取二进制最低位的 1
     * 例如:lowbit(6) = lowbit(110₂) = 2 = 10₂
     */
    private int lowbit(int x) {
        return x & (-x);
    }

    /**
     * 单点更新:在位置 index 增加 delta
     * 时间复杂度:O(log n)
     * 
     * 更新路径:index -> index + lowbit(index) -> ... -> n
     * 每次跳转到父节点,父节点管辖更大的区间
     */
    public void update(int index, long delta) {
        if (index <= 0 || index > n) {
            throw new IllegalArgumentException("Index out of bounds: " + index);
        }

        // 树状数组使用 1-based 索引
        for (int i = index; i <= n; i += lowbit(i)) {
            tree[i] += delta;
        }
    }

    /**
     * 前缀和查询:查询 [1, index] 区间的和
     * 时间复杂度:O(log n)
     * 
     * 查询路径:index -> index - lowbit(index) -> ... -> 0
     * 每次去掉最低位的 1,将前缀分解为若干段
     */
    public long queryPrefix(int index) {
        if (index <= 0) return 0;
        if (index > n) index = n;

        long sum = 0;
        for (int i = index; i > 0; i -= lowbit(i)) {
            sum += tree[i];
        }
        return sum;
    }

    /**
     * 区间和查询:查询 [left, right] 区间的和
     */
    public long queryRange(int left, int right) {
        if (left > right) return 0;
        return queryPrefix(right) - queryPrefix(left - 1);
    }

    /**
     * 单点查询:查询位置 index 的当前值
     */
    public long queryPoint(int index) {
        return queryRange(index, index);
    }

    /**
     * 初始化:根据原数组构建树状数组
     * 时间复杂度:O(n log n),可通过更高效方法优化到 O(n)
     */
    public void build(long[] arr) {
        if (arr.length != n) {
            throw new IllegalArgumentException("Array length mismatch");
        }

        // 方法1:逐个 update,O(n log n)
        for (int i = 0; i < n; i++) {
            update(i + 1, arr[i]);
        }
    }

    /**
     * O(n) 初始化方法
     * 利用 tree[i] = sum(i - lowbit(i) + 1, i) 的关系
     */
    public void buildFast(long[] arr) {
        if (arr.length != n) {
            throw new IllegalArgumentException("Array length mismatch");
        }

        // 先复制原数组
        for (int i = 1; i <= n; i++) {
            tree[i] = arr[i - 1];
        }

        // 再逐层累加:tree[i] 的值需要传递给它的父节点
        for (int i = 1; i <= n; i++) {
            int parent = i + lowbit(i);
            if (parent <= n) {
                tree[parent] += tree[i];
            }
        }
    }

    public int size() {
        return n;
    }

    /**
     * 打印树状数组内部结构(调试用)
     */
    public void printTree() {
        System.out.println("Tree structure:");
        for (int i = 1; i <= n; i++) {
            int lb = lowbit(i);
            System.out.printf("  tree[%d] (lowbit=%d) covers [%d, %d] = %d%n", 
                i, lb, i - lb + 1, i, tree[i]);
        }
    }
}

支持区间更新、单点查询的变体

当需要支持”区间增加、单点查询”时,利用差分数组思想:

public class RangeUpdateFenwickTree {
    private final FenwickTree diffTree;

    public RangeUpdateFenwickTree(int n) {
        this.diffTree = new FenwickTree(n);
    }

    /**
     * 区间增加:将 [left, right] 内所有元素增加 delta
     * 差分数组:diff[left] += delta, diff[right+1] -= delta
     */
    public void rangeUpdate(int left, int right, long delta) {
        diffTree.update(left, delta);
        if (right + 1 <= diffTree.size()) {
            diffTree.update(right + 1, -delta);
        }
    }

    /**
     * 单点查询:位置 index 的实际值 = 差分数组前缀和
     */
    public long pointQuery(int index) {
        return diffTree.queryPrefix(index);
    }
}

支持区间更新、区间查询的变体

同时支持”区间增加、区间求和”需要维护两个树状数组:

public class RangeUpdateQueryFenwickTree {
    private final FenwickTree tree1; // 维护 diff[i]
    private final FenwickTree tree2; // 维护 diff[i] * (i-1)
    private final int n;

    public RangeUpdateQueryFenwickTree(int n) {
        this.n = n;
        this.tree1 = new FenwickTree(n);
        this.tree2 = new FenwickTree(n);
    }

    private void internalUpdate(int index, long delta) {
        tree1.update(index, delta);
        tree2.update(index, delta * (index - 1));
    }

    /**
     * 区间增加:[left, right] 增加 delta
     */
    public void rangeUpdate(int left, int right, long delta) {
        internalUpdate(left, delta);
        if (right + 1 <= n) {
            internalUpdate(right + 1, -delta);
        }
    }

    /**
     * 前缀和查询
     * sum[1..x] = sum(diff[i] * x) - sum(diff[i] * (i-1))
     *           = x * query(tree1, x) - query(tree2, x)
     */
    public long prefixQuery(int index) {
        return index * tree1.queryPrefix(index) - tree2.queryPrefix(index);
    }

    /**
     * 区间和查询
     */
    public long rangeQuery(int left, int right) {
        return prefixQuery(right) - prefixQuery(left - 1);
    }
}

二维树状数组

用于矩阵的区间查询和单点更新:

public class FenwickTree2D {
    private final int n, m;
    private final long[][] tree;

    public 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 update(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 queryPrefix(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;
    }

    /**
     * 子矩阵和查询
     */
    public long queryRange(int x1, int y1, int x2, int y2) {
        return queryPrefix(x2, y2) - queryPrefix(x1 - 1, y2) 
             - queryPrefix(x2, y1 - 1) + queryPrefix(x1 - 1, y1 - 1);
    }
}

经典应用:逆序对统计

public class InversionCount {
    /**
     * 使用树状数组统计逆序对数量
     * 逆序对:i < j 且 a[i] > a[j]
     * 
     * 思路:从右向左遍历,tree.query(a[i]) 得到右边比 a[i] 小的元素个数
     */
    public static long countInversions(int[] arr) {
        // 离散化:将数值压缩到 1~n 范围
        int[] sorted = arr.clone();
        java.util.Arrays.sort(sorted);

        int n = arr.length;
        FenwickTree tree = new FenwickTree(n);
        long inversions = 0;

        for (int i = n - 1; i >= 0; i--) {
            // 找到 arr[i] 在 sorted 中的排名(1-based)
            int rank = lowerBound(sorted, arr[i]) + 1;
            // 查询已经插入的、小于 arr[i] 的元素个数
            inversions += tree.queryPrefix(rank - 1);
            // 将当前元素插入树状数组
            tree.update(rank, 1);
        }

        return inversions;
    }

    private static int lowerBound(int[] arr, int target) {
        int left = 0, right = arr.length;
        while (left < right) {
            int mid = (left + right) >>> 1;
            if (arr[mid] < target) left = mid + 1;
            else right = mid;
        }
        return left;
    }
}

测试与验证

public class FenwickTreeDemo {
    public static void main(String[] args) {
        System.out.println("=== 树状数组基础测试 ===\n");

        // 初始化
        long[] arr = {1, 3, 5, 7, 9, 11};
        FenwickTree ft = new FenwickTree(arr.length);
        ft.buildFast(arr);

        ft.printTree();

        System.out.println("\n前缀和 queryPrefix(3): " + ft.queryPrefix(3)); // 1+3+5 = 9
        System.out.println("区间和 queryRange(2, 5): " + ft.queryRange(2, 5)); // 3+5+7+9 = 24
        System.out.println("单点 queryPoint(4): " + ft.queryPoint(4)); // 7

        // 更新测试
        System.out.println("\n在位置 3 增加 10...");
        ft.update(3, 10);
        System.out.println("更新后前缀和 queryPrefix(3): " + ft.queryPrefix(3)); // 1+3+15 = 19
        System.out.println("更新后区间和 queryRange(2, 5): " + ft.queryRange(2, 5)); // 3+15+7+9 = 34

        // 区间更新测试
        System.out.println("\n=== 区间更新、单点查询测试 ===");
        RangeUpdateFenwickTree ruft = new RangeUpdateFenwickTree(5);
        // 初始数组 [0, 0, 0, 0, 0]
        ruft.rangeUpdate(2, 4, 5); // [0, 5, 5, 5, 0]
        ruft.rangeUpdate(3, 5, 3); // [0, 5, 8, 8, 3]
        System.out.println("pointQuery(3): " + ruft.pointQuery(3)); // 8
        System.out.println("pointQuery(1): " + ruft.pointQuery(1)); // 0
        System.out.println("pointQuery(5): " + ruft.pointQuery(5)); // 3

        // 区间更新区间查询测试
        System.out.println("\n=== 区间更新、区间查询测试 ===");
        RangeUpdateQueryFenwickTree ruqft = new RangeUpdateQueryFenwickTree(5);
        ruqft.rangeUpdate(2, 4, 5); // [0, 5, 5, 5, 0]
        ruqft.rangeUpdate(3, 5, 3); // [0, 5, 8, 8, 3]
        System.out.println("rangeQuery(1, 5): " + ruqft.rangeQuery(1, 5)); // 0+5+8+8+3 = 24
        System.out.println("rangeQuery(2, 4): " + ruqft.rangeQuery(2, 4)); // 5+8+8 = 21

        // 二维测试
        System.out.println("\n=== 二维树状数组测试 ===");
        FenwickTree2D ft2d = new FenwickTree2D(3, 3);
        ft2d.update(1, 1, 1);
        ft2d.update(2, 2, 2);
        ft2d.update(3, 3, 3);
        System.out.println("queryRange(1,1,3,3): " + ft2d.queryRange(1, 1, 3, 3)); // 6
        System.out.println("queryRange(2,2,3,3): " + ft2d.queryRange(2, 2, 3, 3)); // 5

        // 逆序对测试
        System.out.println("\n=== 逆序对统计测试 ===");
        int[] testArr = {5, 2, 6, 1, 3};
        long inv = InversionCount.countInversions(testArr);
        System.out.println("数组: [5, 2, 6, 1, 3]");
        System.out.println("逆序对数量: " + inv); // (5,2), (5,1), (5,3), (2,1), (6,1), (6,3) = 6
    }
}

复杂度对比

数据结构 单点修改 区间查询 区间修改 空间
普通数组 O(1) O(n) O(n) O(n)
前缀和数组 O(n) O(1) O(n) O(n)
树状数组 O(log n) O(log n) O(log n)* O(n)
线段树 O(log n) O(log n) O(log n) O(4n)

*注:区间修改需要维护两个树状数组

与线段树的对比

特性 树状数组 线段树
代码量 少(约30行) 多(约100行)
常数因子 小 较大
支持操作 和、积、异或等 任意可结合操作
区间最值 ❌ 不支持 ✅ 支持
懒标记 需要变体 原生支持

当只需要区间和/积且追求代码简洁时,树状数组是首选;需要区间最值或复杂操作时,选择线段树。

总结

树状数组是算法竞赛和工程实践中最优美的数据结构之一。它通过 lowbit 运算将前缀和分解为 O(log n) 段,以极简的代码实现了高效的区间维护。掌握树状数组不仅能解决大量区间查询问题,更能培养对二进制位运算在算法设计中应用的深刻理解。