轮子 - 树状数组 - 《算法》

参考
实现一
- 单点更新、区间查询如何建树和使用
- 区间更新、单点查询如何建树和使用
实现二
- 区间更新、区间查询如何建树和使用

参考

实现一

可用于以下两种情况

单点更新、区间查询

区间更新、单点查询

public class TreeArray {
  // 树状数组，下标从1开始
  private int[] value;
  // 长度
  private final int length;
  public TreeArray(int n) {
      length = n;
      value = new int[length + 1];
  }
  public TreeArray(int[] arr) {
      length = arr.length;
      value = new int[length + 1];
      for (int i = 1; i <= length; ++i) {
          update(i, arr[i - 1]);
      }
  }
  /**
   * k：x的二进制形式从最低位到最高位连续0的个数
   * 求2^k
   *
   * @param x 非负整数
   * @return 2^k
   */
  private int lowBit(int x) {
      return x & (-x);
  }
  /**
   * 原数组的第i个元素增加k，更新树状数组的值
   *
   * @param i 原数组元素索引位置
   * @param k 更新值
   */
  public void update(int i, int k) {
      while (i <= length) {
          value[i] += k;
          i += lowBit(i);
      }
  }
  /**
   * 求[1, i]这i个数的和
   *
   * @param i 索引位置
   * @return 前i个数的和
   */
  public int getSum(int i) {
      int sum = 0;
      while (i >= 1) {
          sum += value[i];
          i -= lowBit(i);
      }
      return sum;
  }
}

单点更新、区间查询如何建树和使用

int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
TreeArray t = new TreeArray(a);
System.out.println(t.getSum(10) - t.getSum(4));

或者

int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
TreeArray t = new TreeArray(a.length);
for (int i = 1; i <= a.length; ++i) {
    t.update(i, a[i - 1]);
}
System.out.println(t.getSum(10) - t.getSum(4));

区间更新、单点查询如何建树和使用

int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
TreeArray t = new TreeArray(a.length);
t.update(1, a[0]);
for (int i = 2; i <= a.length; ++i) {
    t.update(i, a[i - 1] - a[i - 2]);
}
for (int i = 1; i <= 10; ++i) {
    System.out.print(t.getSum(i) + " ");
}
System.out.println();
// [x,y]区间内的数均加上k
int x = 5;
int y = 7;
int k = 10;
t.update(x, k);
t.update(y + 1, -k);
for (int i = 1; i <= 10; ++i) {
    System.out.print(t.getSum(i) + " ");
}

实现二

可用于以下情况

区间更新、区间查询

public class TreeArray {
  // 树状数组，下标从1开始，差分值D[i]，用于求单点值
  private int[] sum1;
  // 树状数组，下标从1开始，D[i]*(i-1)，用于求前缀和
  private int[] sum2;
  // 长度
  private final int length;
  public TreeArray(int[] arr) {
      length = arr.length;
      sum1 = new int[length + 1];
      sum2 = new int[length + 1];
      update(1, arr[0]);
      for (int i = 2; i <= length; ++i) {
          update(i, arr[i - 1] - arr[i - 2]);
      }
  }
  /**
   * k：x的二进制形式从最低位到最高位连续0的个数
   * 求2^k
   *
   * @param x 非负整数
   * @return 2^k
   */
  private int lowBit(int x) {
      return x & (-x);
  }
  /**
   * 原数组的第i个元素增加k，更新树状数组的值
   *
   * @param i 原数组元素索引位置
   * @param k 更新值
   */
  public void update(int i, int k) {
      int i_ = i;
      while (i <= length) {
          sum1[i] += k;
          sum2[i] += k * (i_ - 1);
          i += lowBit(i);
      }
  }
  /**
   * 求[1, i]这i个数的和
   *
   * @param i 索引位置
   * @return 前i个数的和
   */
  public int getSum(int i) {
      int sum = 0;
      int i_ = i;
      while (i >= 1) {
          sum += i_ * sum1[i] - sum2[i];
          i -= lowBit(i);
      }
      return sum;
  }
  /**
   * 区间更新，[l,r]区间内的数均加上k
   *
   * @param l 左边界
   * @param r 右边界
   * @param k 更新值
   */
  public void sectionUpdate(int l, int r, int k) {
      update(l, k);
      update(r + 1, -k);
  }
}

区间更新、区间查询如何建树和使用

int[] a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
TreeArray t = new TreeArray(a);
System.out.println(t.getSum(10) - t.getSum(4));
// [x,y]区间内的数均加上k
int x = 5;
int y = 7;
int k = 10;
t.sectionUpdate(x, y, k);
System.out.println(t.getSum(10) - t.getSum(4));